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

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

from loguru import logger

from nanobot.utils.helpers import ensure_dir, safe_filename


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

    Stores messages in JSONL format for easy reading and persistence.

    Messages are trimmed after consolidation to keep session size manageable.
    """

    key: str  # 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, **kwargs: Any) -> None:
        """Add a message to the session."""
        msg = {
            "role": role,
            "content": content,
            "timestamp": datetime.now().isoformat(),
            **kwargs
        }
        self.messages.append(msg)
        self.updated_at = datetime.now()

    def add_raw_message(self, msg: dict[str, Any]) -> None:
        """Add a pre-formed message dict to the session, preserving all fields."""
        stored = dict(msg)
        if "timestamp" not in stored:
            stored["timestamp"] = datetime.now().isoformat()
        self.messages.append(stored)
        self.updated_at = datetime.now()

    # Fields that are valid in the Anthropic/OpenAI messages API.
    # Everything else (timestamp, tools_used, etc.) is internal metadata.
    _API_FIELDS = {"role", "content", "tool_calls", "tool_call_id", "name", "reasoning_content"}

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

        The server-side context editing API (clear_tool_uses_20250919) handles
        trimming old tool chains safely at token thresholds, so we send the full
        history and let the server decide what to drop.

        Messages with ``_hidden_sig`` get a ``[HIDDEN:{sig}]`` prefix applied to
        their content so the model knows the user never saw them.  The prefix is
        applied at read time (not stored in content) to preserve prompt-cache
        stability: the same prefixed string is produced every turn.

        Returns:
            List of messages in LLM format (API-relevant fields only).
        """
        out: list[dict[str, Any]] = []
        for m in self.messages:
            msg = {k: v for k, v in m.items() if k in self._API_FIELDS and v is not None}
            sig = m.get("_hidden_sig")
            if sig and isinstance(msg.get("content"), str):
                msg["content"] = f"[HIDDEN:{sig}] {msg['content']}"
            out.append(msg)
        return out
    
    def clear(self) -> None:
        """Clear all messages and reset session to initial state."""
        self.messages = []
        self.updated_at = datetime.now()


class SessionManager:
    """
    Manages conversation sessions.

    Sessions are stored as JSONL files in the sessions directory.
    """

    def __init__(self, workspace: Path):
        self.workspace = workspace
        self.sessions_dir = ensure_dir(self.workspace / "sessions")
        self.legacy_sessions_dir = Path.home() / ".nanobot" / "sessions"
        self._cache: dict[str, Session] = {}
    
    def _get_session_path(self, key: str) -> Path:
        """Get the file path for a session."""
        safe_key = safe_filename(key.replace(":", "_"))
        return self.sessions_dir / f"{safe_key}.jsonl"

    def _get_legacy_session_path(self, key: str) -> Path:
        """Legacy global session path (~/.nanobot/sessions/)."""
        safe_key = safe_filename(key.replace(":", "_"))
        return self.legacy_sessions_dir / f"{safe_key}.jsonl"
    
    def get_or_create(self, key: str) -> Session:
        """
        Get an existing session or create a new one.
        
        Args:
            key: Session key (usually channel:chat_id).
        
        Returns:
            The session.
        """
        if key in self._cache:
            return self._cache[key]
        
        session = self._load(key)
        if session is None:
            session = Session(key=key)
        
        self._cache[key] = session
        return session
    
    def _load(self, key: str) -> Session | None:
        """Load a session from disk."""
        path = self._get_session_path(key)
        if not path.exists():
            legacy_path = self._get_legacy_session_path(key)
            if legacy_path.exists():
                try:
                    shutil.move(str(legacy_path), str(path))
                    logger.info("Migrated session {} from legacy path", key)
                except Exception:
                    logger.exception("Failed to migrate session {}", key)

        if not path.exists():
            return None

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

            with open(path, encoding="utf-8") 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
                        # Ignore legacy last_consolidated field
                    else:
                        messages.append(data)

            return Session(
                key=key,
                messages=messages,
                created_at=created_at or datetime.now(),
                metadata=metadata,
            )
        except Exception as e:
            logger.warning("Failed to load session {}: {}", key, e)
            return None
    
    def save(self, session: Session) -> None:
        """Save a session to disk."""
        path = self._get_session_path(session.key)

        with open(path, "w", encoding="utf-8") as f:
            metadata_line = {
                "_type": "metadata",
                "key": session.key,
                "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")
            for msg in session.messages:
                f.write(json.dumps(msg, ensure_ascii=False) + "\n")

        self._cache[session.key] = session
        self._append_audit(session)

    def _append_audit(self, session: Session) -> None:
        """Append session state to an audit log (append-only, rotated monthly)."""
        now = datetime.now()
        safe_key = safe_filename(session.key.replace(":", "_"))
        audit_path = self.sessions_dir / f"{safe_key}.audit.{now:%Y-%m}.jsonl"
        try:
            with open(audit_path, "a", encoding="utf-8") as f:
                marker = {
                    "_type": "save_marker",
                    "timestamp": now.isoformat(),
                    "message_count": len(session.messages),
                }
                f.write(json.dumps(marker, ensure_ascii=False) + "\n")
                for msg in session.messages:
                    f.write(json.dumps(msg, ensure_ascii=False) + "\n")
        except Exception as e:
            logger.warning("Audit log write failed for {}: {}", session.key, e)
    
    def invalidate(self, key: str) -> None:
        """Remove a session from the in-memory cache."""
        self._cache.pop(key, None)
    
    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:
                # Read just the metadata line
                with open(path, encoding="utf-8") as f:
                    first_line = f.readline().strip()
                    if first_line:
                        data = json.loads(first_line)
                        if data.get("_type") == "metadata":
                            key = data.get("key") or path.stem.replace("_", ":", 1)
                            sessions.append({
                                "key": key,
                                "created_at": data.get("created_at"),
                                "updated_at": data.get("updated_at"),
                                "path": str(path)
                            })
            except Exception:
                continue
        
        return sorted(sessions, key=lambda x: x.get("updated_at", ""), reverse=True)
