"""LLMHosts multi-turn context manager -- tracks conversation state across requests.

When a user has an ongoing conversation and the model switches mid-conversation
(e.g. from a local model to a cloud model due to complexity), this manager
ensures context transfers seamlessly by maintaining condensed message history.
"""

from __future__ import annotations

import hashlib
import json
import logging
import uuid
from datetime import datetime, timezone
from typing import Any

from pydantic import BaseModel, ConfigDict, Field

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Data models
# ---------------------------------------------------------------------------


class Conversation(BaseModel):
    """State of a single ongoing conversation."""

    model_config = ConfigDict(frozen=False)

    id: str
    user: str | None = None
    started_at: datetime
    last_activity: datetime
    turn_count: int = 0
    messages_hash: str = ""  # SHA-256 of current messages for change detection
    current_model: str = ""
    current_backend: str = ""
    total_tokens: int = 0
    model_switches: int = 0  # how many times model changed mid-conversation
    messages: list[dict[str, Any]] = Field(default_factory=list, exclude=True)


# ---------------------------------------------------------------------------
# Context manager
# ---------------------------------------------------------------------------


class ConversationManager:
    """Tracks conversation state across requests for multi-turn routing.

    Pure in-memory -- conversations are lightweight and ephemeral.  The
    :attr:`_ttl` controls how long idle conversations are kept before
    expiration cleanup.

    Thread-safety note: this is designed for single-process async usage
    (one event loop).  No locks needed for coroutine-based concurrency.
    """

    _MAX_TRANSFER_MESSAGES = 20  # max messages to include in a context transfer
    _MAX_TRANSFER_CONTENT_LEN = 2000  # max chars per message in transfer

    def __init__(self, max_conversations: int = 1000, ttl_seconds: int = 3600) -> None:
        self._conversations: dict[str, Conversation] = {}
        self._max = max_conversations
        self._ttl = ttl_seconds

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _now() -> datetime:
        return datetime.now(timezone.utc)

    @staticmethod
    def _hash_messages(messages: list[dict[str, Any]]) -> str:
        """Produce a SHA-256 hash of the message list for change detection."""
        canonical = json.dumps(messages, sort_keys=True, separators=(",", ":"), ensure_ascii=True, default=str)
        return hashlib.sha256(canonical.encode("utf-8")).hexdigest()

    def _evict_oldest_if_full(self) -> None:
        """Evict the oldest conversation by last_activity if at capacity."""
        if len(self._conversations) < self._max:
            return
        # Find and remove the oldest by last_activity
        oldest_id: str | None = None
        oldest_time: datetime | None = None
        for cid, conv in self._conversations.items():
            if oldest_time is None or conv.last_activity < oldest_time:
                oldest_time = conv.last_activity
                oldest_id = cid
        if oldest_id is not None:
            del self._conversations[oldest_id]
            logger.debug("Evicted oldest conversation %s (last_activity=%s)", oldest_id, oldest_time)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def get_or_create(self, conversation_id: str | None = None, user: str | None = None) -> Conversation:
        """Get existing conversation or create a new one.

        Parameters
        ----------
        conversation_id:
            Existing conversation ID.  If *None* or not found, a new
            conversation is created with a generated UUID.
        user:
            Optional user identifier to associate with the conversation.

        Returns
        -------
        Conversation
            The retrieved or newly created conversation.
        """
        if conversation_id and conversation_id in self._conversations:
            conv = self._conversations[conversation_id]
            logger.debug("Retrieved conversation %s (turns=%d)", conversation_id, conv.turn_count)
            return conv

        # Create new
        self._evict_oldest_if_full()
        now = self._now()
        new_id = conversation_id or uuid.uuid4().hex[:16]
        conv = Conversation(
            id=new_id,
            user=user,
            started_at=now,
            last_activity=now,
        )
        self._conversations[new_id] = conv
        logger.debug("Created new conversation %s for user=%s", new_id, user)
        return conv

    def update(
        self,
        conversation_id: str,
        model: str,
        messages: list[dict[str, Any]],
        backend: str,
        tokens: int,
    ) -> None:
        """Update conversation state after a request completes.

        Parameters
        ----------
        conversation_id:
            The conversation to update.
        model:
            Model used for this turn.
        messages:
            Full message list from the request.
        backend:
            Backend that served the request (e.g. "ollama", "openai").
        tokens:
            Total tokens consumed in this turn.
        """
        conv = self._conversations.get(conversation_id)
        if conv is None:
            logger.warning("Cannot update unknown conversation %s", conversation_id)
            return

        # Detect model switch
        if conv.current_model and conv.current_model != model:
            conv.model_switches += 1
            logger.info(
                "Model switch in conversation %s: %s -> %s (switch #%d)",
                conversation_id,
                conv.current_model,
                model,
                conv.model_switches,
            )

        conv.current_model = model
        conv.current_backend = backend
        conv.turn_count += 1
        conv.total_tokens += tokens
        conv.last_activity = self._now()
        conv.messages_hash = self._hash_messages(messages)

        # Store messages for potential context transfer
        conv.messages = list(messages)

        logger.debug(
            "Updated conversation %s: turn=%d model=%s tokens=%d",
            conversation_id,
            conv.turn_count,
            model,
            tokens,
        )

    def get_context_for_transfer(self, conversation_id: str) -> list[dict[str, Any]]:
        """When model switches, extract transferable context.

        Returns a condensed message history suitable for a new model.
        Long messages are truncated, and the total message count is capped.
        System messages are preserved but condensed.

        Parameters
        ----------
        conversation_id:
            The conversation whose context to extract.

        Returns
        -------
        list[dict]
            A condensed list of messages ready to send to a new model.
        """
        conv = self._conversations.get(conversation_id)
        if conv is None:
            logger.warning("Cannot transfer context for unknown conversation %s", conversation_id)
            return []

        if not conv.messages:
            return []

        messages = conv.messages
        max_msgs = self._MAX_TRANSFER_MESSAGES
        max_content = self._MAX_TRANSFER_CONTENT_LEN

        # If too many messages, keep the system message(s) + last N user/assistant turns
        system_msgs = [m for m in messages if m.get("role") == "system"]
        non_system = [m for m in messages if m.get("role") != "system"]

        # Budget: system messages always included, remaining budget for recent turns
        budget = max_msgs - len(system_msgs)
        if budget < 2:
            budget = 2  # always keep at least the last 2 non-system messages

        recent = non_system[-budget:] if len(non_system) > budget else non_system

        transferred: list[dict[str, Any]] = []

        # Add system messages first (condensed)
        for msg in system_msgs:
            transferred.append(self._condense_message(msg, max_content))

        # If we truncated history, add a context note
        if len(non_system) > budget:
            transferred.append(
                {
                    "role": "system",
                    "content": (
                        f"[Context transfer: {len(non_system) - budget} earlier messages omitted. "
                        f"This conversation has {conv.turn_count} turns across "
                        f"{conv.model_switches + 1} model(s).]"
                    ),
                }
            )

        # Add recent messages
        for msg in recent:
            transferred.append(self._condense_message(msg, max_content))

        logger.debug(
            "Context transfer for %s: %d -> %d messages",
            conversation_id,
            len(messages),
            len(transferred),
        )
        return transferred

    @staticmethod
    def _condense_message(msg: dict[str, Any], max_content_len: int) -> dict[str, Any]:
        """Truncate a message's content if it exceeds the limit."""
        result = dict(msg)
        content = result.get("content")

        if isinstance(content, str) and len(content) > max_content_len:
            result["content"] = content[:max_content_len] + f"... [truncated, {len(content)} chars total]"
        elif isinstance(content, list):
            # Multi-modal -- condense text blocks, drop image blocks
            condensed: list[dict[str, Any]] = []
            for block in content:
                if isinstance(block, dict):
                    if block.get("type") == "text":
                        text = str(block.get("text", ""))
                        if len(text) > max_content_len:
                            text = text[:max_content_len] + f"... [truncated, {len(text)} chars total]"
                        condensed.append({"type": "text", "text": text})
                    elif block.get("type") == "image_url":
                        condensed.append({"type": "text", "text": "[image omitted for transfer]"})
                    else:
                        condensed.append(block)
            result["content"] = condensed

        return result

    def cleanup_expired(self) -> int:
        """Remove expired conversations.  Returns count removed."""
        now = self._now()
        expired_ids: list[str] = []

        for cid, conv in self._conversations.items():
            age_seconds = (now - conv.last_activity).total_seconds()
            if age_seconds > self._ttl:
                expired_ids.append(cid)

        for cid in expired_ids:
            del self._conversations[cid]

        if expired_ids:
            logger.info("Cleaned up %d expired conversations", len(expired_ids))
        return len(expired_ids)

    @property
    def active_count(self) -> int:
        """Number of currently tracked conversations."""
        return len(self._conversations)

    def get_conversation(self, conversation_id: str) -> Conversation | None:
        """Retrieve a conversation by ID without creating one."""
        return self._conversations.get(conversation_id)
