"""
session.py
----------
Multi-session continuity for HIEROMEM.

Provides session management and cross-session memory linking,
enabling HIEROMEM to track users across multiple conversations.

Author: Antigravity
"""

from __future__ import annotations

from contextlib import contextmanager
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Dict, Iterator, List, Optional
import uuid


@dataclass
class Session:
    """A user session."""

    session_id: str
    user_id: Optional[str]
    started_at: datetime
    ended_at: Optional[datetime] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

    @property
    def is_active(self) -> bool:
        return self.ended_at is None

    @property
    def duration_seconds(self) -> float:
        end = self.ended_at or datetime.now()
        return (end - self.started_at).total_seconds()

    def end(self) -> None:
        """Mark session as ended."""
        self.ended_at = datetime.now()

    def to_dict(self) -> Dict[str, Any]:
        return {
            "session_id": self.session_id,
            "user_id": self.user_id,
            "started_at": self.started_at.isoformat(),
            "ended_at": self.ended_at.isoformat() if self.ended_at else None,
            "duration_seconds": self.duration_seconds,
            "metadata": self.metadata,
        }


@dataclass
class UserProfile:
    """Aggregated user profile from session history."""

    user_id: str
    total_sessions: int
    total_duration_seconds: float
    first_seen: datetime
    last_seen: datetime
    session_ids: List[str]
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "user_id": self.user_id,
            "total_sessions": self.total_sessions,
            "total_duration_seconds": self.total_duration_seconds,
            "first_seen": self.first_seen.isoformat(),
            "last_seen": self.last_seen.isoformat(),
            "session_ids": self.session_ids,
            "metadata": self.metadata,
        }


class SessionManager:
    """
    Manages user sessions and cross-session memory.

    Example:
        manager = SessionManager(memory)

        with manager.session("user_123"):
            memory.remember("User prefers dark mode")
            # Automatically tagged with user_id and session_id

        # Later, in a new session
        with manager.session("user_123"):
            # Can retrieve past memories from this user
            history = manager.get_user_history("user_123")
    """

    def __init__(self, memory: Any, redis_url: Optional[str] = None):
        """
        Initialize session manager.

        Args:
            memory: HIEROMEM Memory instance
            redis_url: Optional Redis URL for distributed session storage.
                       If None, uses in-memory storage.
        """
        self._memory = memory

        # Avoid circular import
        from hieromem_core.session_store import InMemorySessionStore, RedisSessionStore

        if redis_url:
            self._store = RedisSessionStore(redis_url)
        else:
            self._store = InMemorySessionStore()

        self._current_session: Optional[Session] = None

    @property
    def current_session(self) -> Optional[Session]:
        """Get current active session."""
        return self._current_session

    @property
    def current_user_id(self) -> Optional[str]:
        """Get current user ID."""
        if self._current_session:
            return self._current_session.user_id
        return None

    @property
    def current_session_id(self) -> Optional[str]:
        """Get current session ID."""
        if self._current_session:
            return self._current_session.session_id
        return None

    def start_session(
        self,
        user_id: Optional[str] = None,
        session_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Session:
        """
        Start a new session.

        Args:
            user_id: Optional user identifier
            session_id: Optional session ID (auto-generated if None)
            metadata: Optional session metadata

        Returns:
            The new Session
        """
        session_id = session_id or str(uuid.uuid4())

        session = Session(
            session_id=session_id,
            user_id=user_id,
            started_at=datetime.now(),
            metadata=metadata or {},
        )

        self._store.save_session(session)
        self._current_session = session

        return session

    def end_session(self, session_id: Optional[str] = None) -> Optional[Session]:
        """
        End a session.

        Args:
            session_id: Session to end (current if None)

        Returns:
            The ended Session, or None if not found
        """
        if session_id is None and self._current_session:
            session_id = self._current_session.session_id

        if session_id:
            session = self._store.get_session(session_id)
            if session:
                session.end()
                self._store.save_session(session)
                if (
                    self._current_session
                    and self._current_session.session_id == session_id
                ):
                    self._current_session = None
                return session

        return None

    @contextmanager
    def session(
        self,
        user_id: Optional[str] = None,
        session_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Iterator[Session]:
        """
        Context manager for session scope.

        Args:
            user_id: Optional user identifier
            session_id: Optional session ID
            metadata: Optional session metadata

        Yields:
            The active Session

        Example:
            with manager.session("user_123"):
                memory.remember("Something")  # Tagged with session
        """
        session = self.start_session(user_id, session_id, metadata)
        try:
            yield session
        finally:
            self.end_session(session.session_id)

    def get_session(self, session_id: str) -> Optional[Session]:
        """Get a session by ID."""
        return self._store.get_session(session_id)

    def get_user_sessions(self, user_id: str) -> List[Session]:
        """Get all sessions for a user."""
        return self._store.get_user_sessions(user_id)

    def get_user_profile(self, user_id: str) -> Optional[UserProfile]:
        """
        Get aggregated user profile.

        Args:
            user_id: User identifier

        Returns:
            UserProfile or None if user not found
        """
        sessions = self.get_user_sessions(user_id)
        if not sessions:
            return None

        return UserProfile(
            user_id=user_id,
            total_sessions=len(sessions),
            total_duration_seconds=sum(s.duration_seconds for s in sessions),
            first_seen=min(s.started_at for s in sessions),
            last_seen=max(s.ended_at or s.started_at for s in sessions),
            session_ids=[s.session_id for s in sessions],
        )

    def get_user_history(
        self,
        user_id: str,
        k: int = 20,
        from_session: Optional[str] = None,
    ) -> List[Any]:
        """
        Get memory history for a user.

        Args:
            user_id: User identifier
            k: Max memories to return
            from_session: Optional, only from this session

        Returns:
            List of episodes from user's history
        """
        # Build filter for retrieval
        episodes = []
        for ep in self._memory.episodes:
            if ep.meta.get("user_id") == user_id:
                if from_session and ep.meta.get("session_id") != from_session:
                    continue
                episodes.append(ep)

        # Sort by time (newest first)
        episodes.sort(key=lambda e: e.tocc, reverse=True)
        return episodes[:k]

    def recall_with_session_boost(
        self,
        query: str,
        k: int = 10,
        session_boost: float = 1.5,
        user_boost: float = 1.2,
    ) -> List[Any]:
        """
        Recall memories with boosting for same session/user.

        Memories from the current session are boosted, followed by
        memories from the same user in other sessions.

        Args:
            query: Query string
            k: Number of results
            session_boost: Score multiplier for same session
            user_boost: Score multiplier for same user

        Returns:
            List of boosted episodes
        """
        # Get base retrieval results

        if not hasattr(self._memory, "encoder"):
            return []

        query_emb = self._memory.encoder.encode(query)
        result = self._memory.dual_retrieve(query_vector=query_emb, k=k * 2)

        # Apply boosts
        boosted = []
        for ep in result.episodes:
            score = getattr(ep, "_retrieval_score", 1.0)

            # Check session match
            if self._current_session:
                if ep.meta.get("session_id") == self._current_session.session_id:
                    score *= session_boost
                elif ep.meta.get("user_id") == self._current_session.user_id:
                    score *= user_boost

            boosted.append((score, ep))

        # Sort by boosted score
        boosted.sort(key=lambda x: -x[0])
        return [ep for _, ep in boosted[:k]]
