import time
import json
import logging
from collections import defaultdict, deque
from datetime import datetime, timezone

logger = logging.getLogger("session_manager")

try:
    import redis.asyncio as aioredis
    REDIS_AVAILABLE = True
except ImportError:
    REDIS_AVAILABLE = False

def normalize_keys(obj):
    """Convertit récursivement les clés snake_case du payload OCPP en camelCase."""
    if isinstance(obj, dict):
        new_dict = {}
        for k, v in obj.items():
            parts = k.split('_')
            new_key = parts[0] + ''.join(p.title() for p in parts[1:])
            new_dict[new_key] = normalize_keys(v)
        return new_dict
    elif isinstance(obj, list):
        return [normalize_keys(v) for v in obj]
    return obj

class SessionManager:
    def __init__(self, redis_url=None):
        self.redis = None
        self.memory_cps = {}
        self.memory_frames = defaultdict(lambda: deque(maxlen=200))
        self.pending_calls = {} 
        self.action_map = {} 
        
        if redis_url and REDIS_AVAILABLE:
            try:
                self.redis = aioredis.from_url(redis_url)
                logger.info("SessionManager connecté à Redis.")
            except Exception as e:
                logger.warning(f"Redis indisponible: {e}")
                self.redis = None

    async def register_cp(self, cp_id, model, vendor, subprotocol):
        data = {
            "id": cp_id, "model": str(model), "vendor": str(vendor), 
            "last_activity": time.time(), "status": "Connected", "msg_count": 0,
            "boot_time_ms": None, "reconnections": 0, "clock_drift": None, 
            "rtt_avg_ms": 0, "subprotocol": subprotocol, "last_disconnect_code": None
        }
        if self.redis:
            await self.redis.hset("chargepoints", cp_id, json.dumps(data))
        else:
            self.memory_cps[cp_id] = data

    async def update_disconnect(self, cp_id, code):
        if self.redis:
            data = await self.redis.hget("chargepoints", cp_id)
            if data:
                cp = json.loads(data)
                cp["last_disconnect_code"] = code
                await self.redis.hset("chargepoints", cp_id, json.dumps(cp))
        elif cp_id in self.memory_cps:
            self.memory_cps[cp_id]["last_disconnect_code"] = code

    async def update_cp_status(self, cp_id, status):
        if self.redis:
            data = await self.redis.hget("chargepoints", cp_id)
            if data:
                cp = json.loads(data)
                cp["status"] = str(status)
                cp["last_activity"] = time.time()
                await self.redis.hset("chargepoints", cp_id, json.dumps(cp))
        elif cp_id in self.memory_cps:
            self.memory_cps[cp_id]["status"] = str(status)
            self.memory_cps[cp_id]["last_activity"] = time.time()

    async def purge_frames(self, cp_id):
        if self.redis:
            await self.redis.delete(f"frames:{cp_id}")
            await self.redis.delete(f"stream:{cp_id}")
            data = await self.redis.hget("chargepoints", cp_id)
            if data:
                cp = json.loads(data)
                cp["msg_count"] = 0
                await self.redis.hset("chargepoints", cp_id, json.dumps(cp))
        elif cp_id in self.memory_frames:
            self.memory_frames[cp_id].clear()
            if cp_id in self.memory_cps:
                self.memory_cps[cp_id]["msg_count"] = 0

    async def log_frame(self, cp_id, direction, messageTypeId, action, uniqueId, payload):
        now = time.time()
        rtt_ms = None
        schema_error = None
        
        # Normalisation du payload en camelCase
        payload = normalize_keys(payload)
        
        if messageTypeId == 2: 
            self.action_map[uniqueId] = action
            self.pending_calls[uniqueId] = now
        elif messageTypeId == 3 and action is None: 
            rtt_ms = (now - self.pending_calls.pop(uniqueId, now)) * 1000
            action = self.action_map.pop(uniqueId, "Unknown")
        elif messageTypeId == 4 and action is None: 
            rtt_ms = (now - self.pending_calls.pop(uniqueId, now)) * 1000
            action = self.action_map.pop(uniqueId, "Unknown")

        frame = {
            "timestamp": now, "direction": direction, "messageTypeId": messageTypeId,
            "action": action, "uniqueId": uniqueId, "payload": payload
        }
            
        if self.redis:
            stream_frame = {
                "timestamp": str(now), "direction": str(direction), "messageTypeId": str(messageTypeId),
                "action": str(action), "uniqueId": str(uniqueId), "payload": json.dumps(payload)
            }
            await self.redis.xadd(f"stream:{cp_id}", stream_frame)
            
            await self.redis.lpush(f"frames:{cp_id}", json.dumps(frame))
            await self.redis.ltrim(f"frames:{cp_id}", 0, 199)
            await self.redis.expire(f"frames:{cp_id}", 86400) # TTL 24h
            
            await self.redis.publish("live_frames", json.dumps({"cp_id": cp_id, **frame}))
            
            data = await self.redis.hget("chargepoints", cp_id)
            if data:
                cp = json.loads(data)
                cp["msg_count"] = cp.get("msg_count", 0) + 1
                if action in ["StartTransaction", "StopTransaction"] and "timestamp" in payload:
                    try:
                        cp_time = datetime.fromisoformat(payload["timestamp"].replace("Z", "+00:00"))
                        cp["clock_drift"] = abs((datetime.now(timezone.utc) - cp_time).total_seconds())
                    except: pass
                if rtt_ms: cp["rtt_avg_ms"] = int(rtt_ms)
                await self.redis.hset("chargepoints", cp_id, json.dumps(cp))
        else:
            self.memory_frames[cp_id].append(frame)
            if cp_id in self.memory_cps:
                self.memory_cps[cp_id]["msg_count"] += 1

    async def get_cps(self):
        if self.redis:
            data = await self.redis.hgetall("chargepoints")
            return [json.loads(v) for v in data.values()]
        return list(self.memory_cps.values())

    async def get_frames(self, cp_id):
        if self.redis:
            data = await self.redis.lrange(f"frames:{cp_id}", 0, -1)
            return [json.loads(f) for f in reversed(data)]
        return list(self.memory_frames[cp_id])

session_mgr = SessionManager(redis_url="redis://redis:6379")