from __future__ import annotations

import importlib
import json
from dataclasses import dataclass, field
from typing import Any
from typing import Protocol


@dataclass
class SessionState:
    session_id: str
    user_id: str
    channel: str
    status: str = "idle"
    current_intent: str | None = None
    pending_slots: list[str] = field(default_factory=list)
    slots: dict[str, Any] = field(default_factory=dict)
    workflow: dict[str, Any] | None = None
    routing_debug: dict[str, Any] | None = None
    last_user_text: str | None = None
    last_agent_text: str | None = None
    context_memory: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "session_id": self.session_id,
            "user_id": self.user_id,
            "channel": self.channel,
            "status": self.status,
            "current_intent": self.current_intent,
            "pending_slots": self.pending_slots,
            "slots": self.slots,
            "workflow": self.workflow,
            "routing_debug": self.routing_debug,
            "last_user_text": self.last_user_text,
            "last_agent_text": self.last_agent_text,
            "context_memory": self.context_memory,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "SessionState":
        return cls(
            session_id=data["session_id"],
            user_id=data["user_id"],
            channel=data.get("channel", "app"),
            status=data.get("status", "idle"),
            current_intent=data.get("current_intent"),
            pending_slots=list(data.get("pending_slots", [])),
            slots=dict(data.get("slots", {})),
            workflow=data.get("workflow"),
            routing_debug=data.get("routing_debug"),
            last_user_text=data.get("last_user_text"),
            last_agent_text=data.get("last_agent_text"),
            context_memory=dict(data.get("context_memory", {})),
        )


class SessionStore(Protocol):
    def get_or_create(self, session_id: str, user_id: str, channel: str = "app") -> SessionState:
        ...

    def get(self, session_id: str) -> SessionState | None:
        ...

    def save(self, session: SessionState) -> SessionState:
        ...


class InMemorySessionStore:
    def __init__(self) -> None:
        self._sessions: dict[str, SessionState] = {}

    def get_or_create(self, session_id: str, user_id: str, channel: str = "app") -> SessionState:
        session = self._sessions.get(session_id)
        if session is None:
            session = SessionState(session_id=session_id, user_id=user_id, channel=channel)
            self._sessions[session_id] = session
        return session

    def get(self, session_id: str) -> SessionState | None:
        return self._sessions.get(session_id)

    def save(self, session: SessionState) -> SessionState:
        self._sessions[session.session_id] = session
        return session


class RedisSessionStore:
    def __init__(
        self,
        redis_url: str,
        key_prefix: str = "agent:session",
        ttl_seconds: int = 86400,
    ) -> None:
        redis_module = importlib.import_module("redis")
        self._client = redis_module.Redis.from_url(redis_url, decode_responses=True)
        self._key_prefix = key_prefix
        self._ttl_seconds = ttl_seconds

    def get_or_create(self, session_id: str, user_id: str, channel: str = "app") -> SessionState:
        session = self.get(session_id)
        if session is not None:
            return session

        session = SessionState(session_id=session_id, user_id=user_id, channel=channel)
        self.save(session)
        return session

    def get(self, session_id: str) -> SessionState | None:
        payload = self._client.get(self._build_key(session_id))
        if payload is None:
            return None
        return SessionState.from_dict(json.loads(payload))

    def save(self, session: SessionState) -> SessionState:
        self._client.set(
            self._build_key(session.session_id),
            json.dumps(session.to_dict(), ensure_ascii=False),
            ex=self._ttl_seconds,
        )
        return session

    def _build_key(self, session_id: str) -> str:
        return f"{self._key_prefix}:{session_id}"