"""Session management and context window handling."""

from __future__ import annotations

import json
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any

import tiktoken
from database import Database
from logging_config import get_logger

log = get_logger(__name__)


@dataclass
class Session:
    """A conversation session with message history."""

    id: str
    title: str
    created_at: str
    updated_at: str
    messages: list[dict[str, Any]]
    metadata: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def new(cls, system_prompt: str, title: str = "") -> Session:
        """Create a new session with a system message."""
        now = datetime.now(timezone.utc).isoformat()
        return cls(
            id=uuid.uuid4().hex[:12],
            title=title,
            created_at=now,
            updated_at=now,
            messages=[{"role": "system", "content": system_prompt}],
        )

    def touch(self) -> None:
        """Update the last-modified timestamp."""
        self.updated_at = datetime.now(timezone.utc).isoformat()

    def auto_title(self) -> None:
        """Set title from first user message if not already set."""
        if self.title:
            return
        for msg in self.messages:
            if msg["role"] == "user":
                text = msg["content"]
                self.title = text[:60] + ("..." if len(text) > 60 else "")
                return

    def to_dict(self) -> dict[str, Any]:
        """Serialize session to a plain dict."""
        return {
            "id": self.id,
            "title": self.title,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
            "messages": self.messages,
            "metadata": self.metadata,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> Session:
        """Deserialize a session from a plain dict."""
        return cls(
            id=data["id"],
            title=data.get("title", ""),
            created_at=data["created_at"],
            updated_at=data["updated_at"],
            messages=data["messages"],
            metadata=data.get("metadata", {}),
        )


class SessionStore:
    """SQLite-backed session persistence."""

    def __init__(self, db: Database) -> None:
        self._db = db
        self._log = log.bind(component="session_store")

    def save(self, session: Session) -> None:
        """Save (upsert) a session to the database."""
        session.touch()
        self._db.conn.execute(
            """INSERT OR REPLACE INTO sessions
               (id, title, created_at, updated_at, messages, metadata)
               VALUES (?, ?, ?, ?, ?, ?)""",
            (
                session.id,
                session.title,
                session.created_at,
                session.updated_at,
                json.dumps(session.messages, ensure_ascii=False),
                json.dumps(session.metadata, ensure_ascii=False),
            ),
        )
        self._db.conn.commit()
        self._log.debug(
            "session_saved",
            session_id=session.id,
            message_count=len(session.messages),
        )

    def load(self, session_id: str) -> Session:
        """Load a session from the database."""
        row = self._db.conn.execute("SELECT * FROM sessions WHERE id = ?", (session_id,)).fetchone()
        if row is None:
            raise FileNotFoundError(f"Session not found: {session_id}")
        self._log.debug("session_loaded", session_id=session_id)
        return Session.from_dict(
            {
                "id": row["id"],
                "title": row["title"],
                "created_at": row["created_at"],
                "updated_at": row["updated_at"],
                "messages": json.loads(row["messages"]),
                "metadata": json.loads(row["metadata"]),
            }
        )

    def list_sessions(self, limit: int = 20) -> list[dict[str, Any]]:
        """List recent sessions (summary only, sorted newest first)."""
        rows = self._db.conn.execute(
            """SELECT id, title, updated_at, messages
               FROM sessions ORDER BY updated_at DESC LIMIT ?""",
            (limit,),
        ).fetchall()
        return [
            {
                "id": row["id"],
                "title": row["title"] or "(untitled)",
                "updated_at": row["updated_at"],
                "message_count": len(json.loads(row["messages"])),
            }
            for row in rows
        ]

    def search_sessions(self, query: str, limit: int = 100) -> list[dict[str, Any]]:
        """Search sessions by ID, title, and message content."""
        normalized = query.strip().lower()
        if not normalized:
            return self.list_sessions(limit=limit)

        pattern = f"%{normalized}%"
        rows = self._db.conn.execute(
            """SELECT id, title, updated_at, messages
               FROM sessions
               WHERE LOWER(id) LIKE ?
                  OR LOWER(title) LIKE ?
                  OR LOWER(messages) LIKE ?
               ORDER BY updated_at DESC
               LIMIT ?""",
            (pattern, pattern, pattern, limit),
        ).fetchall()

        return [
            {
                "id": row["id"],
                "title": row["title"] or "(untitled)",
                "updated_at": row["updated_at"],
                "message_count": (len(json.loads(row["messages"])) if row["messages"] else 0),
            }
            for row in rows
        ]

    def delete(self, session_id: str) -> bool:
        """Delete a session. Returns True if deleted, False if not found."""
        cursor = self._db.conn.execute("DELETE FROM sessions WHERE id = ?", (session_id,))
        self._db.conn.commit()
        if cursor.rowcount > 0:
            self._log.info("session_deleted", session_id=session_id)
            return True
        return False


class ContextManager:
    """Manages message truncation to fit within a model's context window."""

    def __init__(
        self,
        max_tokens: int = 120_000,
        reserve_tokens: int = 4_000,
        encoding_name: str = "cl100k_base",
    ) -> None:
        self.max_tokens = max_tokens
        self.reserve_tokens = reserve_tokens
        self._encoding = tiktoken.get_encoding(encoding_name)
        self._log = log.bind(component="context_manager")

    def count_tokens(self, text: str) -> int:
        """Count tokens in a text string."""
        return len(self._encoding.encode(text))

    def _message_tokens(self, message: dict[str, Any]) -> int:
        """Estimate token count for a single message."""
        tokens = 4  # per-message overhead (role, separators)
        content = message.get("content")
        if isinstance(content, str):
            tokens += self.count_tokens(content)
        elif isinstance(content, list):
            # Multimodal content (text + image parts)
            for part in content:
                if isinstance(part, dict):
                    if part.get("type") == "text":
                        tokens += self.count_tokens(part.get("text", ""))
                    elif part.get("type") == "image_url":
                        # Each image costs ~1000-2000 tokens depending on
                        # resolution; use a conservative estimate.
                        tokens += 2000
        tool_calls = message.get("tool_calls")
        if tool_calls:
            tokens += self.count_tokens(json.dumps(tool_calls))
        return tokens

    @staticmethod
    def _group_into_turns(
        messages: list[dict[str, Any]],
    ) -> list[list[dict[str, Any]]]:
        """Group messages into conversational turns.

        Each turn starts with a user message and includes all subsequent
        assistant/tool messages until the next user message.  This ensures
        tool-call groups are never split.
        """
        turns: list[list[dict[str, Any]]] = []
        current: list[dict[str, Any]] = []

        for msg in messages:
            if msg["role"] == "user" and current:
                turns.append(current)
                current = []
            current.append(msg)

        if current:
            turns.append(current)
        return turns

    def prepare_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """Return a (possibly truncated) copy that fits the context window.

        The system message is always preserved.  Oldest turns are dropped
        first.  Tool-call groups are never split.
        """
        if not messages:
            return []

        budget = self.max_tokens - self.reserve_tokens

        # System message is always kept
        system_msg = messages[0] if messages[0]["role"] == "system" else None
        rest = messages[1:] if system_msg else messages

        system_tokens = self._message_tokens(system_msg) if system_msg else 0
        turns = self._group_into_turns(rest)
        turn_tokens = [sum(self._message_tokens(m) for m in turn) for turn in turns]

        # Keep turns from the end until the budget is exhausted
        remaining = budget - system_tokens
        kept: list[list[dict[str, Any]]] = []

        for turn, tokens in reversed(list(zip(turns, turn_tokens))):
            if remaining - tokens < 0 and kept:
                break
            kept.append(turn)
            remaining -= tokens

        kept.reverse()

        dropped = len(turns) - len(kept)
        if dropped > 0:
            self._log.info(
                "context_truncated",
                turns_total=len(turns),
                turns_kept=len(kept),
                turns_dropped=dropped,
            )

        result: list[dict[str, Any]] = []
        if system_msg:
            result.append(system_msg)
        for turn in kept:
            result.extend(turn)
        return result
