"""Conversation memory with SQLite persistence."""

from __future__ import annotations

import logging
import uuid
from datetime import UTC, datetime

import aiosqlite

from agntspace.llm.base import Message

log = logging.getLogger(__name__)


class ConversationMemory:
    """Manages conversation history in SQLite."""

    def __init__(self, db: aiosqlite.Connection) -> None:
        self._db = db
        self._skill_loader: object | None = None  # post-wired from main.py
        # Per-instance migration flag — each new ConversationMemory (e.g.
        # one per test) runs its own additive migration. Previously this
        # was a class-level flag which broke fresh in-memory test DBs.
        self._migrations_applied: bool = False

    async def _ensure_schema(self) -> None:
        """Apply additive column migrations once per connection.

        ``streaming INTEGER DEFAULT 0`` marks a message whose content is
        still being written (draft assistant replies during a chat turn).
        The UI reads this flag to render an in-flight cursor and to show
        partial text on reconnect. Pre-migration rows default to 0 which
        is the correct "finalized" semantics.
        """
        if self._migrations_applied:
            return
        try:
            await self._db.execute(
                "ALTER TABLE messages ADD COLUMN streaming INTEGER DEFAULT 0"
            )
            await self._db.commit()
        except Exception:
            # Column already exists — idempotent
            pass
        self._migrations_applied = True

    def _load_auto_title_config(self) -> dict:
        """Load auto-title config from the conversation meta-skill."""
        if not self._skill_loader:
            return {}
        try:
            cfg = self._skill_loader.load_skill_config_file("conversation", "auto-title.yaml")
            return cfg if isinstance(cfg, dict) else {}
        except Exception:
            return {}

    async def create_conversation(
        self, title: str | None = None, conversation_id: str | None = None
    ) -> str:
        """Create a new conversation. Returns conversation ID."""
        cid = conversation_id or uuid.uuid4().hex
        now = datetime.now(UTC).isoformat()
        await self._db.execute(
            "INSERT INTO conversations (id, title, created_at, updated_at) VALUES (?, ?, ?, ?)",
            (cid, title or "New conversation", now, now),
        )
        await self._db.commit()
        return cid

    async def add_message(
        self, conversation_id: str, role: str, content: str,
        provenance: dict | None = None,
    ) -> str:
        """Store a message. Returns message ID.

        Lazily creates the conversation row if it does not exist yet, so
        opening a chat UI without sending anything never leaves an empty
        row in the list.

        *provenance* is an optional dict stamped on user messages at the
        channel boundary — ``{kind, source_channel, origin_session,
        source_tool}``.  Stored as JSON in the ``provenance`` column so
        every message row carries its full origin context.
        """
        import json as _json

        mid = uuid.uuid4().hex
        now = datetime.now(UTC).isoformat()
        token_count = len(content) // 4  # rough approximation
        prov_json = _json.dumps(provenance) if provenance else None
        await self._db.execute(
            "INSERT OR IGNORE INTO conversations (id, title, created_at, updated_at) "
            "VALUES (?, ?, ?, ?)",
            (conversation_id, "New conversation", now, now),
        )
        await self._db.execute(
            "INSERT INTO messages (id, conversation_id, role, content, created_at, token_count, provenance) "
            "VALUES (?, ?, ?, ?, ?, ?, ?)",
            (mid, conversation_id, role, content, now, token_count, prov_json),
        )
        # Touch updated_at and auto-unarchive: a new message on an archived
        # thread pulls it back into the sidebar, matching every messaging app.
        await self._db.execute(
            "UPDATE conversations SET updated_at = ?, archived = 0 WHERE id = ?",
            (now, conversation_id),
        )
        await self._db.commit()
        return mid

    async def upsert_assistant_message(
        self,
        conversation_id: str,
        content: str,
        *,
        message_id: str | None = None,
        streaming: bool = True,
    ) -> str:
        """Create-or-update an assistant message. Returns the message id.

        The incremental-persist path (2026-04-19): chat_stream writes a
        draft assistant row at stream start (no ``message_id``), then
        updates the same row as chunks arrive (with ``message_id`` +
        ``streaming=True``), then finalizes on ``done=True`` (with
        ``message_id`` + ``streaming=False``). If the WebSocket is
        dropped mid-stream, whatever content has been persisted so far
        survives — the user sees partial text on reconnect instead of
        an empty page.

        Lazily creates the conversation row if missing, same as
        ``add_message``. Updates ``conversations.updated_at`` on every
        call so the conversation sorts correctly in the sidebar.
        """
        await self._ensure_schema()
        now = datetime.now(UTC).isoformat()
        streaming_flag = 1 if streaming else 0
        token_count = len(content) // 4

        if message_id is None:
            mid = uuid.uuid4().hex
            await self._db.execute(
                "INSERT OR IGNORE INTO conversations (id, title, created_at, updated_at) "
                "VALUES (?, ?, ?, ?)",
                (conversation_id, "New conversation", now, now),
            )
            await self._db.execute(
                "INSERT INTO messages "
                "(id, conversation_id, role, content, created_at, token_count, streaming) "
                "VALUES (?, ?, ?, ?, ?, ?, ?)",
                (mid, conversation_id, "assistant", content, now, token_count, streaming_flag),
            )
            await self._db.execute(
                "UPDATE conversations SET updated_at = ?, archived = 0 WHERE id = ?",
                (now, conversation_id),
            )
            await self._db.commit()
            return mid

        # UPDATE path — same row, new content. Do NOT rewrite created_at
        # (clients key off it for ordering). Refresh conversation
        # updated_at so the sidebar surfaces the activity.
        await self._db.execute(
            "UPDATE messages SET content = ?, token_count = ?, streaming = ? WHERE id = ?",
            (content, token_count, streaming_flag, message_id),
        )
        await self._db.execute(
            "UPDATE conversations SET updated_at = ?, archived = 0 WHERE id = ?",
            (now, conversation_id),
        )
        await self._db.commit()
        return message_id

    async def add_tool_message(
        self,
        conversation_id: str,
        tool_name: str,
        tool_input: dict,
        result: dict,
        *,
        status: str = "ok",
        duration_ms: int = 0,
    ) -> str:
        """Persist a tool-call trace as a ``role=tool`` message.

        Every tool call the agent runs during a chat turn leaves a
        visible trace in conversation history so the UI can render it
        as an expandable card ("Claude used tool X → {result preview}").
        Content is a compact JSON payload: the tool name, key input
        hints, outcome classification, truncated result preview.
        Returns the message id so downstream handlers can cross-link.
        """
        import json as _json

        await self._ensure_schema()
        payload = {
            "tool": tool_name,
            "input_keys": sorted(tool_input.keys()) if isinstance(tool_input, dict) else [],
            "status": status,
            "duration_ms": duration_ms,
        }
        if isinstance(result, dict):
            payload["result_keys"] = sorted(result.keys())[:12]
            err = result.get("error")
            if isinstance(err, str) and err:
                payload["error_preview"] = err[:240]
        body = _json.dumps(payload)

        mid = uuid.uuid4().hex
        now = datetime.now(UTC).isoformat()
        await self._db.execute(
            "INSERT OR IGNORE INTO conversations (id, title, created_at, updated_at) "
            "VALUES (?, ?, ?, ?)",
            (conversation_id, "New conversation", now, now),
        )
        await self._db.execute(
            "INSERT INTO messages "
            "(id, conversation_id, role, content, created_at, token_count, streaming) "
            "VALUES (?, ?, ?, ?, ?, ?, 0)",
            (mid, conversation_id, "tool", body, now, len(body) // 4),
        )
        await self._db.execute(
            "UPDATE conversations SET updated_at = ?, archived = 0 WHERE id = ?",
            (now, conversation_id),
        )
        await self._db.commit()
        return mid

    async def get_messages(
        self, conversation_id: str, limit: int = 50
    ) -> list[Message]:
        """Get recent messages for a conversation, ordered by time."""
        cursor = await self._db.execute(
            "SELECT role, content FROM messages "
            "WHERE conversation_id = ? ORDER BY created_at ASC LIMIT ?",
            (conversation_id, limit),
        )
        rows = await cursor.fetchall()
        return [Message(role=row[0], content=row[1]) for row in rows]

    async def get_messages_detailed(
        self, conversation_id: str, limit: int = 100,
    ) -> list[dict]:
        """Fetch messages with the metadata the UI needs to render them:
        id, role, content, streaming flag, created_at. The streaming flag
        lets the chat UI show an in-flight cursor on draft assistant rows
        (set to True while a turn is being written, cleared when the turn
        finalizes). Used by ``GET /api/conversations/{id}/messages``.
        """
        await self._ensure_schema()
        cursor = await self._db.execute(
            "SELECT id, role, content, created_at, COALESCE(streaming, 0) "
            "FROM messages WHERE conversation_id = ? "
            "ORDER BY created_at ASC LIMIT ?",
            (conversation_id, limit),
        )
        rows = await cursor.fetchall()
        return [
            {
                "id": row[0],
                "role": row[1],
                "content": row[2],
                "created_at": row[3],
                "streaming": bool(row[4]),
            }
            for row in rows
        ]

    async def get_message(
        self, conversation_id: str, message_id: str
    ) -> dict | None:
        """Fetch a single message by id, scoped to its conversation.

        Used by the edit/regenerate flow to validate that the target message
        belongs to the conversation the user is editing in. Returns None if
        the message does not exist OR does not belong to this conversation.
        """
        cursor = await self._db.execute(
            "SELECT id, role, content, created_at FROM messages "
            "WHERE id = ? AND conversation_id = ?",
            (message_id, conversation_id),
        )
        row = await cursor.fetchone()
        if not row:
            return None
        return {
            "id": row[0],
            "role": row[1],
            "content": row[2],
            "created_at": row[3],
        }

    async def update_user_message_content(
        self, conversation_id: str, message_id: str, new_content: str
    ) -> bool:
        """Rewrite a user message's content. Returns True if updated.

        Refuses to touch non-user messages — agents should not be able to
        rewrite their own past output, and the UI never offers it. The
        caller (PATCH route) is responsible for truncating subsequent
        messages so the conversation stays coherent.
        """
        cursor = await self._db.execute(
            "SELECT role FROM messages WHERE id = ? AND conversation_id = ?",
            (message_id, conversation_id),
        )
        row = await cursor.fetchone()
        if not row or row[0] != "user":
            return False
        token_count = len(new_content) // 4
        await self._db.execute(
            "UPDATE messages SET content = ?, token_count = ? WHERE id = ?",
            (new_content, token_count, message_id),
        )
        now = datetime.now(UTC).isoformat()
        await self._db.execute(
            "UPDATE conversations SET updated_at = ? WHERE id = ?",
            (now, conversation_id),
        )
        await self._db.commit()
        return True

    async def delete_messages_after(
        self,
        conversation_id: str,
        message_id: str,
        *,
        inclusive: bool = False,
    ) -> int:
        """Delete every message in this conversation that comes after
        ``message_id`` in time order. Returns the deleted row count.

        Used by:
          - **edit**: ``inclusive=False`` — keep the edited user message,
            wipe the assistant reply (and any tool traces) that came from
            the previous version of it.
          - **regenerate**: ``inclusive=False`` against the last user
            message — same shape, wipe everything after the user's last
            turn.

        Cutoff is the target message's ``created_at``. We delete strictly
        after when ``inclusive=False``, otherwise from-and-including the
        cutoff. Returns ``0`` if the target message does not exist (no-op
        rather than raise — caller may have raced with another delete).
        """
        cursor = await self._db.execute(
            "SELECT created_at FROM messages WHERE id = ? AND conversation_id = ?",
            (message_id, conversation_id),
        )
        row = await cursor.fetchone()
        if not row:
            return 0
        cutoff = row[0]
        op = ">=" if inclusive else ">"
        cursor = await self._db.execute(
            f"DELETE FROM messages WHERE conversation_id = ? AND created_at {op} ?",
            (conversation_id, cutoff),
        )
        deleted = cursor.rowcount or 0
        # Refresh updated_at so the sidebar order stays current.
        now = datetime.now(UTC).isoformat()
        await self._db.execute(
            "UPDATE conversations SET updated_at = ? WHERE id = ?",
            (now, conversation_id),
        )
        await self._db.commit()
        return deleted

    async def get_last_user_message(self, conversation_id: str) -> dict | None:
        """Return the most recent ``role=user`` message, or None.

        Used by the regenerate flow to find what the agent should reply
        to after the assistant turn has been wiped.
        """
        cursor = await self._db.execute(
            "SELECT id, role, content, created_at FROM messages "
            "WHERE conversation_id = ? AND role = 'user' "
            "ORDER BY created_at DESC LIMIT 1",
            (conversation_id,),
        )
        row = await cursor.fetchone()
        if not row:
            return None
        return {
            "id": row[0],
            "role": row[1],
            "content": row[2],
            "created_at": row[3],
        }

    async def get_messages_since(
        self,
        since_iso: str,
        limit: int = 500,
    ) -> list[dict]:
        """Return all messages newer than ``since_iso`` across every conversation.

        Phase 6 (memory dreaming) needs this to feed the light dream phase
        with recent chat experience. Each row is a dict with the
        conversation id, role, content, and created_at timestamp — enough
        for the dream gathering helper to pair user→assistant turns and
        produce stable episode_ids.

        Returns rows in chronological order across all conversations
        (oldest first). The hard cap is a safety net against runaway
        dream phases on a heavily-used vault; in normal use the lookback
        window keeps the result well under it.
        """
        cursor = await self._db.execute(
            "SELECT conversation_id, id, role, content, created_at FROM messages "
            "WHERE created_at >= ? ORDER BY created_at ASC LIMIT ?",
            (since_iso, limit),
        )
        rows = await cursor.fetchall()
        return [
            {
                "conversation_id": row[0],
                "message_id": row[1],
                "role": row[2],
                "content": row[3],
                "created_at": row[4],
            }
            for row in rows
        ]

    async def get_context_messages(
        self, conversation_id: str, max_tokens: int = 8000
    ) -> list[Message]:
        """Get messages that fit within a token budget.

        Takes the most recent messages up to max_tokens.
        Token count approximated as len(content) / 4.
        Uses SQL LIMIT to avoid fetching entire conversation history.

        **Only returns ``user`` and ``assistant`` rows.** ``tool`` rows
        exist in the DB so the UI can render tool traces in history, but
        they are NOT replayed to the LLM on subsequent turns — OpenAI's
        strict format requires every ``tool`` message to follow an
        ``assistant`` message whose ``tool_calls`` array references it by
        id, and we don't reconstruct that scaffolding. The final assistant
        reply already summarized the tool outcome, which is sufficient LLM
        context. Filtering here (not at the call site) guarantees the
        invariant holds for every future caller.
        """
        # Fetch at most 100 recent messages — far more than any token budget needs
        cursor = await self._db.execute(
            "SELECT role, content, token_count FROM messages "
            "WHERE conversation_id = ? AND role IN ('user', 'assistant') "
            "ORDER BY created_at DESC LIMIT 100",
            (conversation_id,),
        )
        rows = await cursor.fetchall()

        selected: list[Message] = []
        total_tokens = 0
        for row in rows:
            tokens = row[2] or len(row[1]) // 4
            if total_tokens + tokens > max_tokens:
                break
            selected.append(Message(role=row[0], content=row[1]))
            total_tokens += tokens

        selected.reverse()  # oldest first
        return selected

    async def list_conversations(self, limit: int = 0, include_archived: bool = False) -> list[dict]:
        """List recent conversations with metadata.

        Only conversations that have at least one message are returned —
        empty rows (from WS connects that never sent a message, or legacy
        auto-created rows) are hidden.  ``limit=0`` (default) returns all.
        """
        base = (
            "SELECT c.id, c.title, c.created_at, c.updated_at, c.archived "
            "FROM conversations c "
            "WHERE EXISTS (SELECT 1 FROM messages m WHERE m.conversation_id = c.id)"
        )
        if include_archived:
            query = base + " ORDER BY c.updated_at DESC"
        else:
            query = base + " AND c.archived = 0 ORDER BY c.updated_at DESC"
        params: tuple = ()
        if limit > 0:
            query += " LIMIT ?"
            params = (limit,)
        cursor = await self._db.execute(query, params)
        rows = await cursor.fetchall()
        return [
            {
                "id": row[0],
                "title": row[1],
                "created_at": row[2],
                "updated_at": row[3],
                "archived": bool(row[4]) if len(row) > 4 else False,
            }
            for row in rows
        ]

    async def archive_conversation(self, conversation_id: str) -> None:
        """Archive a conversation (hidden from default list, not deleted)."""
        await self._db.execute(
            "UPDATE conversations SET archived = 1 WHERE id = ?",
            (conversation_id,),
        )
        await self._db.commit()

    async def unarchive_conversation(self, conversation_id: str) -> None:
        """Restore an archived conversation."""
        await self._db.execute(
            "UPDATE conversations SET archived = 0 WHERE id = ?",
            (conversation_id,),
        )
        await self._db.commit()

    async def delete_conversation(self, conversation_id: str) -> None:
        """Permanently delete a conversation and all its messages."""
        await self._db.execute("DELETE FROM messages WHERE conversation_id = ?", (conversation_id,))
        await self._db.execute("DELETE FROM conversations WHERE id = ?", (conversation_id,))
        await self._db.commit()

    async def update_title(self, conversation_id: str, title: str) -> None:
        """Update a conversation's title."""
        await self._db.execute(
            "UPDATE conversations SET title = ? WHERE id = ?",
            (title, conversation_id),
        )
        await self._db.commit()

    async def auto_title(self, conversation_id: str, llm=None) -> str | None:
        """Auto-generate a title from the conversation using LLM.

        Only updates if current title is 'New conversation'.
        Uses the LLM to create a short descriptive title from the
        first user message and assistant response.
        Returns the new title or None if not changed.
        """
        cursor = await self._db.execute(
            "SELECT title FROM conversations WHERE id = ?", (conversation_id,),
        )
        row = await cursor.fetchone()
        if not row or row[0] != "New conversation":
            return None

        # Grab first few messages for context
        cursor = await self._db.execute(
            "SELECT role, content FROM messages WHERE conversation_id = ? "
            "ORDER BY created_at ASC LIMIT 4",
            (conversation_id,),
        )
        rows = await cursor.fetchall()
        if not rows:
            return None

        # Build a snippet of the conversation for the LLM
        _title_cfg = self._load_auto_title_config()
        _snippet_len = _title_cfg.get("snippet_length", 300)
        snippet_parts = []
        for role, content in rows:
            text = content.strip()[:_snippet_len]
            snippet_parts.append(f"{role}: {text}")
        snippet = "\n".join(snippet_parts)

        title = None

        # Load auto-title config from conversation skill YAML
        _title_cfg = self._load_auto_title_config()
        _sys_prompt = _title_cfg.get("system_prompt", (
            "Generate a short conversation title (max 6 words) that captures "
            "the topic of this conversation. Return ONLY the title text, "
            "nothing else. No quotes, no punctuation at the end, no prefix."
        ))
        _max_tokens = _title_cfg.get("max_tokens", 30)
        _temperature = _title_cfg.get("temperature", 0.3)
        _max_title_len = _title_cfg.get("max_title_length", 80)

        if llm:
            try:
                response = await llm.complete(
                    messages=[Message(role="user", content=snippet)],
                    system=_sys_prompt.strip(),
                    max_tokens=_max_tokens,
                    temperature=_temperature,
                )
                title = response.content.strip().strip('"\'').strip()
                log.debug("LLM auto-title: %r", title)
                # Sanity check — if LLM returned something too long or empty, fall back
                if not title or len(title) > _max_title_len:
                    title = None
            except Exception:
                log.exception("Auto-title LLM call failed")  # Don't swallow silently

        if not title:
            # Fallback: truncate first user message
            user_text = rows[0][1].strip()
            title = user_text[:50].rstrip()
            if len(user_text) > 50:
                title += "…"
            title = title.replace("\n", " ")

        if not title:
            return None

        await self.update_title(conversation_id, title)
        return title

    async def purge_empty(self) -> int:
        """Delete conversation rows that have no messages. Returns count deleted."""
        cursor = await self._db.execute(
            "DELETE FROM conversations WHERE NOT EXISTS "
            "(SELECT 1 FROM messages m WHERE m.conversation_id = conversations.id)"
        )
        await self._db.commit()
        return cursor.rowcount or 0

    async def conversation_exists(self, conversation_id: str) -> bool:
        """Check if a conversation exists."""
        cursor = await self._db.execute(
            "SELECT 1 FROM conversations WHERE id = ?", (conversation_id,)
        )
        return await cursor.fetchone() is not None

    # ── Artifact context (feedback loop) ──

    async def set_artifact(
        self,
        conversation_id: str,
        artifact_type: str,
        artifact_path: str,
        metadata: dict | None = None,
    ) -> None:
        """Set the in-focus artifact for a conversation.

        Called when a tool produces a mutable output (presentation, article,
        report). The artifact stays in focus until replaced or cleared, so
        follow-up messages like "make slide 3 more detailed" know what to
        refine.
        """
        import json
        meta_json = json.dumps(metadata) if metadata else None
        try:
            await self._db.execute(
                "UPDATE conversations SET artifact_type = ?, artifact_path = ?, "
                "artifact_meta = ? WHERE id = ?",
                (artifact_type, artifact_path, meta_json, conversation_id),
            )
            await self._db.commit()
        except Exception:
            log.debug("Failed to set artifact for %s", conversation_id, exc_info=True)

    async def get_artifact(self, conversation_id: str) -> dict | None:
        """Get the in-focus artifact for a conversation, if any."""
        import json
        try:
            cursor = await self._db.execute(
                "SELECT artifact_type, artifact_path, artifact_meta "
                "FROM conversations WHERE id = ?",
                (conversation_id,),
            )
            row = await cursor.fetchone()
            if row and row[0]:
                return {
                    "type": row[0],
                    "path": row[1],
                    "metadata": json.loads(row[2]) if row[2] else {},
                }
        except Exception:
            log.debug("Failed to get artifact for %s", conversation_id, exc_info=True)
        return None

    async def clear_artifact(self, conversation_id: str) -> None:
        """Clear the in-focus artifact for a conversation."""
        try:
            await self._db.execute(
                "UPDATE conversations SET artifact_type = NULL, artifact_path = NULL, "
                "artifact_meta = NULL WHERE id = ?",
                (conversation_id,),
            )
            await self._db.commit()
        except Exception:
            pass
