"""Conversation history adapter for pydantic-ai agents.

Cache strategy (write-through, per-process):
  - ``load_conversation_history``: cache-first; on miss loads from DB and
    stores the result so subsequent calls in the same process never hit DB again.
  - ``append_to_history_cache``: called after every ``save_message``; appends
    the new message directly to the in-memory list so the cache grows with the
    conversation.  Creates the entry if not yet present (covers new chats where
    ``load_conversation_history`` may not have been called yet).
  - ``serialize_history`` / ``deserialize_history``: compact wire format
    ``[{role, content}]`` so rt_agent can pass its already-loaded history to
    tatasteel and ui_agent in the HTTP request body — those processes have their
    own caches and would otherwise each make a redundant DB round-trip.
"""
from __future__ import annotations

import json

from loguru import logger as _log
from pydantic_ai.messages import ModelMessage, ModelRequest, ModelResponse, TextPart, UserPromptPart

from backend.chanakya.conversation.database import get_messages

_HISTORY_CACHE: dict[int, list[ModelMessage]] = {}


# ---------------------------------------------------------------------------
# Wire-format helpers — used by rt_agent to pass history to sub-agents
# ---------------------------------------------------------------------------

def serialize_history(history: list[ModelMessage]) -> list[dict[str, str]]:
    """Flatten ModelMessages to ``[{role, content}]`` for HTTP transport."""
    result: list[dict[str, str]] = []
    for msg in history:
        if isinstance(msg, ModelRequest):
            for part in msg.parts:
                if hasattr(part, "content"):
                    result.append({"role": "user", "content": str(part.content)})
        elif isinstance(msg, ModelResponse):
            for part in msg.parts:
                if hasattr(part, "content"):
                    result.append({"role": "assistant", "content": str(part.content)})
    return result


def deserialize_history(data: list[dict[str, str]]) -> list[ModelMessage]:
    """Reconstruct ModelMessages from ``[{role, content}]`` transport format."""
    result: list[ModelMessage] = []
    for item in data:
        role = item.get("role", "")
        content = item.get("content", "")
        if role == "user":
            result.append(ModelRequest(parts=[UserPromptPart(content=content)]))
        elif role == "assistant":
            result.append(ModelResponse(parts=[TextPart(content=content)]))
    return result


# ---------------------------------------------------------------------------
# Write-through cache helpers
# ---------------------------------------------------------------------------

def append_to_history_cache(
    conversation_id: int | None,
    author: str,
    message: str,
    tatasteel_cache: dict | None = None,
) -> None:
    """Append a newly saved message to the in-memory cache.

    Creates the cache entry when it does not yet exist — this covers new
    conversations where ``load_conversation_history`` may not have been called
    before the first ``save_message``.

    Args:
        conversation_id: Conversation the message belongs to.
        author:          ``"user"`` or ``"assistant"``.
        message:         Plain-text message content.
        tatasteel_cache: For assistant messages, ``{rows, explanation, question}``
                         from metadata — rows are embedded so visualization
                         follow-ups work without a DB re-query.
    """
    if not conversation_id:
        return

    if conversation_id not in _HISTORY_CACHE:
        _HISTORY_CACHE[conversation_id] = []

    if author == "user":
        _HISTORY_CACHE[conversation_id].append(
            ModelRequest(parts=[UserPromptPart(content=message)])
        )
    elif author == "assistant":
        content = message
        if tatasteel_cache and tatasteel_cache.get("rows"):
            content += (
                f"\n\n[Previous query data"
                f" — explanation: {tatasteel_cache.get('explanation', '')}"
                f"\nrows: {json.dumps(tatasteel_cache['rows'])}]"
            )
        _HISTORY_CACHE[conversation_id].append(
            ModelResponse(parts=[TextPart(content=content)])
        )

    _log.debug(
        "[HISTORY] Appended {} message to cache for conversation {} (total={})",
        author, conversation_id, len(_HISTORY_CACHE[conversation_id]),
    )


# ---------------------------------------------------------------------------
# Load
# ---------------------------------------------------------------------------

async def load_conversation_history(conversation_id: int | None) -> list[ModelMessage]:
    """Return conversation history as a pydantic-ai message list.

    Cache-first: if the conversation is already in memory, returns that copy
    without touching the DB.  Falls back to a DB load and caches the result.

    Args:
        conversation_id: The conversation to load, or ``None`` / ``0`` to skip.

    Returns:
        Ordered ``list[ModelMessage]`` for ``agent.run(message_history=...)``.
        Returns ``[]`` when *conversation_id* is falsy or on any DB error.
    """
    if not conversation_id:
        return []

    cached = _HISTORY_CACHE.get(conversation_id)
    if cached is not None:
        _log.debug(
            "[HISTORY] Cache hit for conversation {} ({} messages)",
            conversation_id, len(cached),
        )
        return cached

    try:
        messages = await get_messages(conversation_id)
        history: list[ModelMessage] = []
        for msg in messages:
            if msg["message_author"] == "user":
                history.append(ModelRequest(parts=[UserPromptPart(content=msg["message"])]))
            elif msg["message_author"] == "assistant":
                content = msg["message"]
                raw_meta = msg.get("metadata")
                if raw_meta:
                    try:
                        metadata = json.loads(raw_meta) if isinstance(raw_meta, str) else raw_meta
                        cache = metadata.get("_tatasteel_cache")
                        if cache and cache.get("rows"):
                            content += (
                                f"\n\n[Previous query data"
                                f" — explanation: {cache.get('explanation', '')}"
                                f"\nrows: {json.dumps(cache['rows'])}]"
                            )
                    except Exception:
                        pass
                history.append(ModelResponse(parts=[TextPart(content=content)]))
        _log.debug(
            "[HISTORY] Loaded {} messages from DB for conversation {}",
            len(history), conversation_id,
        )
        _HISTORY_CACHE[conversation_id] = history
        return history
    except Exception as exc:
        _log.warning(
            "[HISTORY] Failed to load conversation history for {}: {}",
            conversation_id, exc,
        )
        return []