"""Search index for thread messages.

Provides a lightweight full-text search index that lives alongside the
LangGraph checkpoint tables in ``~/.icarus/sessions.db``.  Messages are
indexed after each turn so the desktop search modal can query them
without deserializing checkpoint blobs.
"""

from __future__ import annotations

import logging
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING, Literal, NotRequired, TypedDict

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    import aiosqlite

from icarus.session.threads import (
    _connect,
    _get_custom_titles,
    _table_exists,
    list_threads,
)

_MAX_CONTENT_LENGTH = 2000
_SNIPPET_RADIUS = 50


class SearchResult(TypedDict):
    """A single search result returned by :func:`search_threads`."""

    thread_id: str
    title: str
    updated_at: str | None
    message_count: int
    match_type: Literal["title", "content"]
    snippet: str | None
    status: NotRequired[str]  # added by sidecar enrichment


class SearchPage(TypedDict):
    """Paginated search response returned by :func:`search_threads`."""

    results: list[SearchResult]
    next_cursor: tuple[str, str] | None


_search_table_ensured: set[str] = set()


async def _ensure_search_index_table(conn: aiosqlite.Connection, *, db_key: str = "_default") -> None:
    """Create the ``message_search_index`` table if it doesn't exist.

    Idempotent — safe to call on every write path.  Uses a module-level
    set keyed by database identity so ``CREATE TABLE`` runs once per
    distinct database, not once per process.  This is critical in
    multi-user server mode where each user has their own SQLite file.

    Args:
        conn: An open aiosqlite connection.
        db_key: Identifier for the database (e.g. the path string).
            Defaults to ``"_default"`` for single-user desktop/CLI.
    """
    if db_key in _search_table_ensured:
        return
    await conn.execute("""
        CREATE TABLE IF NOT EXISTS message_search_index (
            thread_id  TEXT NOT NULL,
            message_id TEXT NOT NULL,
            role       TEXT NOT NULL,
            content    TEXT NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            PRIMARY KEY (thread_id, message_id)
        )
    """)
    _search_table_ensured.add(db_key)


async def index_thread_messages(
    thread_id: str,
    human_text: str,
    ai_text: str,
    *,
    db_path: Path | None = None,
) -> None:
    """Index a human/AI message pair for search.

    Inserts two rows (one per role) into the ``message_search_index`` table.
    Content is truncated to ``_MAX_CONTENT_LENGTH`` characters to keep the
    index compact.

    Args:
        thread_id: The thread these messages belong to.
        human_text: The user's message content.
        ai_text: The assistant's response content.
        db_path: Explicit path to the SQLite database.  When ``None``
            (the default) falls back to :func:`~icarus.session.threads.get_db_path`.
    """
    now = datetime.now(timezone.utc).isoformat()

    async with _connect(db_path=db_path) as conn:
        await _ensure_search_index_table(conn, db_key=str(db_path or "_default"))

        for role, text in (("human", human_text), ("ai", ai_text)):
            message_id = str(uuid.uuid4())
            truncated = text[:_MAX_CONTENT_LENGTH]
            await conn.execute(
                """
                INSERT INTO message_search_index
                    (thread_id, message_id, role, content, created_at)
                VALUES (?, ?, ?, ?, ?)
                """,
                (thread_id, message_id, role, truncated, now),
            )

        await conn.commit()


async def safe_index_thread(
    thread_id: str,
    human_content: str,
    ai_text: str,
    *,
    db_path: Path | None = None,
) -> None:
    """Fire-and-forget wrapper around :func:`index_thread_messages`.

    Swallows all exceptions so callers can schedule this without worrying
    about indexing errors bubbling up into stream handling.
    """
    try:
        await index_thread_messages(thread_id, human_content, ai_text, db_path=db_path)
    except Exception:
        logger.debug("Search indexing failed for %s", thread_id, exc_info=True)


def _extract_snippet(content: str, query: str) -> str | None:
    """Extract ~100 chars around the first match of *query* in *content*.

    Uses a radius of ``_SNIPPET_RADIUS`` chars on each side.  Adds ``...``
    prefix/suffix when the snippet is truncated.

    Args:
        content: The full text to search in.
        query: The substring to find (case-insensitive).

    Returns:
        The snippet string, or ``None`` if *query* is not found.
    """
    lower_content = content.lower()
    lower_query = query.lower()
    idx = lower_content.find(lower_query)
    if idx == -1:
        return None

    start = max(0, idx - _SNIPPET_RADIUS)
    end = min(len(content), idx + len(query) + _SNIPPET_RADIUS)

    snippet = content[start:end]

    if start > 0:
        snippet = "..." + snippet
    if end < len(content):
        snippet = snippet + "..."

    return snippet


async def delete_search_index(thread_id: str, *, db_path: Path | None = None) -> None:
    """Delete all search index entries for a thread.

    Args:
        thread_id: The thread whose index entries should be removed.
        db_path: Explicit path to the SQLite database.  When ``None``
            (the default) falls back to :func:`~icarus.session.threads.get_db_path`.
    """
    async with _connect(db_path=db_path) as conn:
        await _ensure_search_index_table(conn, db_key=str(db_path or "_default"))
        await conn.execute(
            "DELETE FROM message_search_index WHERE thread_id = ?",
            (thread_id,),
        )
        await conn.commit()


async def search_threads(
    query: str,
    limit: int = 20,
    cursor: tuple[str, str] | None = None,
    time_range: tuple[str, str] | None = None,
    *,
    db_path: Path | None = None,
) -> SearchPage:
    """Search threads by title and message content.

    Args:
        query: Search query string.  Empty string returns recent threads.
        limit: Maximum number of results to return.
        cursor: Optional ``(updated_at, thread_id)`` cursor for pagination.
        time_range: Optional ``(start_iso, end_iso)`` tuple to filter by
            ``updated_at``.
        db_path: Explicit path to the SQLite database.  When ``None``
            (the default) falls back to :func:`~icarus.session.threads.get_db_path`.

    Returns:
        Dict with ``results`` (list of :class:`SearchResult`) and
        ``next_cursor`` (tuple or None) for pagination.
    """
    if not query.strip():
        threads = await list_threads(
            limit=limit, cursor=cursor, include_message_count=True,
            db_path=db_path,
        )
        results = [
            SearchResult(
                thread_id=t["thread_id"],
                title=t.get("title", t["thread_id"][:8] + "..."),
                updated_at=t.get("updated_at"),
                message_count=t.get("message_count", 0),
                match_type="title",
                snippet=None,
            )
            for t in threads
        ]
        next_cursor = None
        if threads and len(threads) == limit:
            last = threads[-1]
            next_cursor = (last.get("updated_at", ""), last["thread_id"])
        return {"results": results, "next_cursor": next_cursor}

    results_by_id: dict[str, SearchResult] = {}

    async with _connect(db_path=db_path) as conn:
        await _ensure_search_index_table(conn, db_key=str(db_path or "_default"))

        # Fresh per-user DBs may not have checkpoints yet (created on first stream)
        if not await _table_exists(conn, "checkpoints"):
            return {"results": [], "next_cursor": None}

        escaped = query.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
        like_pattern = f"%{escaped}%"

        if await _table_exists(conn, "thread_metadata"):
            title_query = """
                SELECT m.thread_id, m.custom_title,
                       MAX(json_extract(c.metadata, '$.updated_at')) as updated_at
                FROM thread_metadata m
                LEFT JOIN checkpoints c ON c.thread_id = m.thread_id
                WHERE m.custom_title LIKE ? ESCAPE '\\'
                GROUP BY m.thread_id
                ORDER BY updated_at DESC
            """
            async with conn.execute(title_query, (like_pattern,)) as cur:
                for row in await cur.fetchall():
                    tid, title, updated_at = row
                    results_by_id[tid] = SearchResult(
                        thread_id=tid,
                        title=title,
                        updated_at=updated_at,
                        message_count=0,
                        match_type="title",
                        snippet=None,
                    )

        content_query = """
            SELECT s.thread_id, s.content,
                   MAX(json_extract(c.metadata, '$.updated_at')) as updated_at
            FROM message_search_index s
            LEFT JOIN checkpoints c ON c.thread_id = s.thread_id
            WHERE s.content LIKE ? ESCAPE '\\'
            GROUP BY s.thread_id
            ORDER BY updated_at DESC
        """
        async with conn.execute(content_query, (like_pattern,)) as cur:
            for row in await cur.fetchall():
                tid, content, updated_at = row
                # Dedup: title match wins over content match
                if tid in results_by_id:
                    continue
                snippet = _extract_snippet(content, query)
                results_by_id[tid] = SearchResult(
                    thread_id=tid,
                    title=tid[:8] + "...",  # placeholder, filled below
                    updated_at=updated_at,
                    message_count=0,
                    match_type="content",
                    snippet=snippet,
                )

        # Fill custom titles for content-only matches
        content_only_ids = [
            tid for tid, r in results_by_id.items() if r["match_type"] == "content"
        ]
        if content_only_ids:
            custom_titles = await _get_custom_titles(conn, content_only_ids)
            for tid in content_only_ids:
                custom = custom_titles.get(tid)
                if custom:
                    results_by_id[tid]["title"] = custom

        if results_by_id:
            all_ids = list(results_by_id.keys())
            placeholders = ",".join("?" for _ in all_ids)
            count_query = f"""
                SELECT thread_id, COUNT(*) as cnt
                FROM message_search_index
                WHERE thread_id IN ({placeholders})
                GROUP BY thread_id
            """  # noqa: S608
            async with conn.execute(count_query, all_ids) as cur:
                for row in await cur.fetchall():
                    tid, cnt = row
                    if tid in results_by_id:
                        results_by_id[tid]["message_count"] = cnt

    results = sorted(
        results_by_id.values(),
        key=lambda r: r.get("updated_at") or "",
        reverse=True,
    )

    if time_range:
        start, end = time_range
        results = [
            r for r in results
            if r.get("updated_at") and start <= r["updated_at"] <= end
        ]

    if cursor:
        cursor_updated_at = cursor[0] or ""
        cursor_thread_id = cursor[1] or ""
        filtered = []
        past_cursor = False
        for r in results:
            if not past_cursor:
                if (r.get("updated_at") or "") < cursor_updated_at or (
                    (r.get("updated_at") or "") == cursor_updated_at
                    and r["thread_id"] < cursor_thread_id
                ):
                    past_cursor = True
                else:
                    continue
            if past_cursor:
                filtered.append(r)
        results = filtered

    page = results[:limit]
    next_cursor = None
    if len(page) == limit and len(results) > limit:
        last = page[-1]
        next_cursor = (last.get("updated_at", ""), last["thread_id"])
    return {"results": page, "next_cursor": next_cursor}


async def reindex_all_threads(*, db_path: Path | None = None) -> int:
    """One-time migration: index all existing threads for search.

    Args:
        db_path: Explicit path to the SQLite database.  When ``None``
            (the default) falls back to :func:`~icarus.session.threads.get_db_path`.

    Returns the number of threads indexed.
    """
    from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer

    serde = JsonPlusSerializer()
    count = 0

    threads = await list_threads(limit=9999, db_path=db_path)

    async with _connect(db_path=db_path) as conn:
        await _ensure_search_index_table(conn, db_key=str(db_path or "_default"))

        for thread in threads:
            tid = thread["thread_id"]
            try:
                async with conn.execute(
                    """SELECT type, checkpoint FROM checkpoints
                       WHERE thread_id = ?
                       ORDER BY rowid DESC LIMIT 1""",
                    (tid,),
                ) as cur:
                    row = await cur.fetchone()

                if not row:
                    continue

                data = serde.loads_typed((row[0], row[1]))
                messages = data.get("channel_values", {}).get("messages", [])

                for msg in messages:
                    role = getattr(msg, "type", None)
                    if role not in ("human", "ai"):
                        continue

                    content = getattr(msg, "content", "")
                    if isinstance(content, list):
                        content = " ".join(
                            b.get("text", "") for b in content
                            if isinstance(b, dict) and b.get("type") == "text"
                        )

                    if not content:
                        continue

                    msg_id = getattr(msg, "id", None) or f"{role}-{uuid.uuid4().hex[:8]}"
                    await conn.execute(
                        """INSERT OR REPLACE INTO message_search_index
                           (thread_id, message_id, role, content, created_at)
                           VALUES (?, ?, ?, ?, ?)""",
                        (tid, msg_id, role, content[:_MAX_CONTENT_LENGTH],
                         datetime.now(timezone.utc).isoformat()),
                    )

                count += 1
            except Exception:
                logger.warning("Failed to index thread %s", tid, exc_info=True)

        await conn.commit()

    return count


if __name__ == "__main__":
    import asyncio
    import sys

    logging.basicConfig(level=logging.INFO)

    if "--reindex" in sys.argv:
        indexed = asyncio.run(reindex_all_threads())
        print(f"Indexed {indexed} threads.")
    else:
        print("Usage: python -m icarus.session.search --reindex")
