"""Session management for conversation history."""

import asyncio
import json
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any

from loguru import logger

from vikingbot.config.schema import SessionKey
from vikingbot.sandbox.manager import SandboxManager
from vikingbot.utils.helpers import ensure_dir


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

    Stores messages in JSONL format for easy reading and persistence.
    """

    key: SessionKey  # channel:chat_id
    messages: list[dict[str, Any]] = field(default_factory=list)
    created_at: datetime = field(default_factory=datetime.now)
    updated_at: datetime = field(default_factory=datetime.now)
    metadata: dict[str, Any] = field(default_factory=dict)

    def add_message(
        self, role: str, content: str, sender_id: str | None = None, token_usage: dict[str, Any] = None, **kwargs: Any
    ) -> None:
        """Add a message to the session."""
        msg = {"role": role, "content": content, "timestamp": datetime.now().isoformat(), **kwargs}
        if sender_id is not None:
            msg["sender_id"] = sender_id
        if token_usage is not None:
            msg["token_usage"] = token_usage
        self.messages.append(msg)
        self.updated_at = datetime.now()

    def get_history(self, max_messages: int = 50) -> list[dict[str, Any]]:
        """
        Get message history for LLM context.

        Args:
            max_messages: Maximum messages to return.

        Returns:
            List of messages in LLM format.
        """
        # Get recent messages
        recent = (
            self.messages[-max_messages:] if len(self.messages) > max_messages else self.messages
        )

        # Convert to LLM format (just role and content)
        return [{"role": m["role"], "content": m["content"]} for m in recent]

    def clear(self) -> None:
        """Clear all messages in the session."""
        self.messages = []
        self.updated_at = datetime.now()

    def clone(self) -> "Session":
        """Create a deep copy of this session."""
        import copy

        return Session(
            key=self.key,
            messages=copy.deepcopy(self.messages),
            created_at=self.created_at,
            updated_at=self.updated_at,
            metadata=copy.deepcopy(self.metadata),
        )


class SessionManager:
    """
    Manages conversation sessions with persistence and caching.

    SessionManager handles the lifecycle of conversation sessions, including
    creation, retrieval, caching, and persistent storage. Sessions are stored
    as JSONL files in a designated directory for durability.

    The manager maintains an in-memory cache of active sessions to improve
    performance and reduce disk I/O. Sessions are automatically persisted when
    modified.

    Attributes:
        bot_data_path: Path to the bot's data directory.
        workspace: Path to the workspace directory within bot_data.
        sessions_dir: Path where session JSONL files are stored.
        _cache: In-memory cache mapping SessionKey to Session objects.
        sandbox_manager: Optional sandbox manager for isolated operations.

    Example:
        >>> manager = SessionManager(Path("/path/to/bot/data"))
        >>> session_key = SessionKey(channel="telegram", chat_id="12345")
        >>> session = manager.get_or_create(session_key)
        >>> session.add_message("user", "Hello!")
        >>> await manager.save(session)
    """

    def __init__(
        self,
        bot_data_path: Path,
        sandbox_manager: "SandboxManager | None" = None,
    ):
        self.bot_data_path = bot_data_path
        self.workspace = bot_data_path / "workspace"
        self.sessions_dir = ensure_dir(bot_data_path / "sessions")
        self._cache: dict[SessionKey, Session] = {}
        self.sandbox_manager = sandbox_manager

    def _get_session_path(self, session_key: SessionKey) -> Path:
        return self.sessions_dir / f"{session_key.safe_name()}.jsonl"

    def get_or_create(self, key: SessionKey, skip_heartbeat: bool = False) -> Session:
        """
        Get an existing session or create a new one.

        Args:
            key: Session key (usually channel:chat_id).
            skip_heartbeat: Whether to skip heartbeat for this session.

        Returns:
            The session.
        """
        # Check cache
        if key in self._cache:
            return self._cache[key]

        # Try to load from disk
        session = self._load(key)
        if session is None:
            session = Session(key=key)
            if skip_heartbeat:
                session.metadata["skip_heartbeat"] = True

        self._cache[key] = session

        if self.sandbox_manager:
            from vikingbot.utils.helpers import ensure_session_workspace

            if self.sandbox_manager.config.mode == "shared":
                workspace_path = self.sandbox_manager.workspace / "shared"
            else:
                workspace_path = self.sandbox_manager.workspace / key.safe_name()
            ensure_session_workspace(workspace_path)

        # Initialize sandbox
        if self.sandbox_manager:
            asyncio.create_task(self._init_sandbox(key))

        return session

    async def _init_sandbox(self, key: SessionKey) -> None:
        """Initialize sandbox for a session."""
        if self.sandbox_manager is None:
            return
        try:
            await self.sandbox_manager.get_sandbox(key)
        except Exception as e:
            logger.warning(f"Failed to initialize sandbox for {key}: {e}")

    def _load(self, session_key: SessionKey) -> Session | None:
        """Load a session from disk."""
        path = self._get_session_path(session_key)

        if not path.exists():
            return None

        try:
            messages = []
            metadata = {}
            created_at = None
            session_key_from_metadata = None

            with open(path) as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue

                    data = json.loads(line)

                    if data.get("_type") == "metadata":
                        metadata = data.get("metadata", {})
                        created_at = (
                            datetime.fromisoformat(data["created_at"])
                            if data.get("created_at")
                            else None
                        )
                        session_key_from_metadata = SessionKey.from_safe_name(
                            data.get("session_key")
                        )
                    else:
                        messages.append(data)

            effective_key = session_key_from_metadata if session_key_from_metadata else session_key

            return Session(
                key=effective_key,
                messages=messages,
                created_at=created_at or datetime.now(),
                metadata=metadata,
            )
        except Exception as e:
            logger.warning(f"Failed to load session {session_key}: {e}")
            return None

    async def save(self, session: Session) -> None:
        """Save a session to disk."""
        path = self._get_session_path(session.key)

        with open(path, "w") as f:
            # Write metadata first
            metadata_line = {
                "_type": "metadata",
                "session_key": session.key.safe_name(),
                "created_at": session.created_at.isoformat(),
                "updated_at": session.updated_at.isoformat(),
                "metadata": session.metadata,
            }
            f.write(json.dumps(metadata_line, ensure_ascii=False) + "\n")

            # Write messages
            for msg in session.messages:
                f.write(json.dumps(msg, ensure_ascii=False) + "\n")

        self._cache[session.key] = session

    def delete(self, key: SessionKey) -> bool:
        """
        Delete a session.

        Args:
            key: Session key.

        Returns:
            True if deleted, False if not found.
        """
        # Clean up sandbox if enabled
        if self.sandbox_manager is not None:
            asyncio.create_task(self.sandbox_manager.cleanup_session(key))

        # Remove from cache
        self._cache.pop(key, None)

        # Remove file
        path = self._get_session_path(key)
        if path.exists():
            path.unlink()
            return True
        return False

    def list_sessions(self) -> list[dict[str, Any]]:
        """
        List all sessions.

        Returns:
            List of session info dicts.
        """
        sessions = []

        for path in self.sessions_dir.glob("*.jsonl"):
            try:
                with open(path) as f:
                    first_line = f.readline().strip()
                    if first_line:
                        data = json.loads(first_line)
                        if data.get("_type") == "metadata":
                            session_key = SessionKey.from_safe_name(data.get("session_key"))
                            metadata = data.get("metadata", {})
                            sessions.append(
                                {
                                    "key": session_key,
                                    "created_at": data.get("created_at"),
                                    "updated_at": data.get("updated_at"),
                                    "metadata": metadata,
                                    "path": str(path),
                                }
                            )
            except Exception:
                continue

        return sorted(sessions, key=lambda x: x.get("updated_at", ""), reverse=True)
