"""
session_store.py
----------------
Pluggable session storage backend.

Provides an abstract interface for session storage and concrete implementations
for in-memory and Redis-backed storage.
"""

from __future__ import annotations

import json
from abc import ABC, abstractmethod
from typing import Dict, List, Optional
from datetime import datetime

from hieromem_core.session import Session


class SessionStore(ABC):
    """Abstract interface for session storage."""

    @abstractmethod
    def save_session(self, session: Session) -> None:
        """Save or update a session."""
        pass

    @abstractmethod
    def get_session(self, session_id: str) -> Optional[Session]:
        """Retrieve a session by ID."""
        pass

    @abstractmethod
    def get_user_sessions(self, user_id: str) -> List[Session]:
        """Retrieve all sessions for a user."""
        pass

    @abstractmethod
    def delete_session(self, session_id: str) -> None:
        """Delete a session."""
        pass


class InMemorySessionStore(SessionStore):
    """Reference implementation using in-memory dictionary."""

    def __init__(self):
        self._sessions: Dict[str, Session] = {}
        self._user_sessions: Dict[str, List[str]] = {}

    def save_session(self, session: Session) -> None:
        self._sessions[session.session_id] = session
        if session.user_id:
            if session.user_id not in self._user_sessions:
                self._user_sessions[session.user_id] = []
            if session.session_id not in self._user_sessions[session.user_id]:
                self._user_sessions[session.user_id].append(session.session_id)

    def get_session(self, session_id: str) -> Optional[Session]:
        return self._sessions.get(session_id)

    def get_user_sessions(self, user_id: str) -> List[Session]:
        session_ids = self._user_sessions.get(user_id, [])
        return [self._sessions[sid] for sid in session_ids if sid in self._sessions]

    def delete_session(self, session_id: str) -> None:
        if session_id in self._sessions:
            session = self._sessions[session_id]
            if session.user_id and session.user_id in self._user_sessions:
                self._user_sessions[session.user_id].remove(session_id)
            del self._sessions[session_id]


class RedisSessionStore(SessionStore):
    """Production implementation using Redis."""

    def __init__(self, redis_url: str):
        import redis

        self._redis = redis.from_url(redis_url, decode_responses=True)
        self._ttl = 86400 * 30  # 30 days retention

    def _serialize(self, session: Session) -> str:
        return json.dumps(session.to_dict())

    def _deserialize(self, data: str) -> Session:
        d = json.loads(data)
        return Session(
            session_id=d["session_id"],
            user_id=d["user_id"],
            started_at=datetime.fromisoformat(d["started_at"]),
            ended_at=(
                datetime.fromisoformat(d["ended_at"]) if d.get("ended_at") else None
            ),
            metadata=d.get("metadata", {}),
        )

    def save_session(self, session: Session) -> None:
        key = f"session:{session.session_id}"
        self._redis.setex(key, self._ttl, self._serialize(session))

        if session.user_id:
            user_key = f"user_sessions:{session.user_id}"
            self._redis.sadd(user_key, session.session_id)
            self._redis.expire(user_key, self._ttl)

    def get_session(self, session_id: str) -> Optional[Session]:
        data = self._redis.get(f"session:{session_id}")
        if data:
            return self._deserialize(data)
        return None

    def get_user_sessions(self, user_id: str) -> List[Session]:
        user_key = f"user_sessions:{user_id}"
        session_ids = self._redis.smembers(user_key)
        sessions = []
        for sid in session_ids:
            session = self.get_session(sid)
            if session:
                sessions.append(session)
        # Sort by started_at desc
        sessions.sort(key=lambda s: s.started_at, reverse=True)
        return sessions

    def delete_session(self, session_id: str) -> None:
        session = self.get_session(session_id)
        if session:
            self._redis.delete(f"session:{session_id}")
            if session.user_id:
                self._redis.srem(f"user_sessions:{session.user_id}", session_id)
