"""Thread management using LangGraph's built-in checkpoint persistence."""

from __future__ import annotations

import logging
from contextlib import asynccontextmanager
from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING, NotRequired, TypedDict

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    import aiosqlite
    from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
    from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver

logger = logging.getLogger(__name__)

_aiosqlite_patched = False


def _patch_aiosqlite() -> None:
    """Patch aiosqlite.Connection with `is_alive()` if missing.

    Required by langgraph-checkpoint>=2.1.0.
    See: https://github.com/langchain-ai/langgraph/issues/6583
    """
    global _aiosqlite_patched  # noqa: PLW0603  # Module-level flag requires global statement
    if _aiosqlite_patched:
        return

    import aiosqlite as _aiosqlite

    if not hasattr(_aiosqlite.Connection, "is_alive"):

        def _is_alive(self: _aiosqlite.Connection) -> bool:
            """Check if the connection is still alive.

            Returns:
                True if connection is alive, False otherwise.
            """
            return bool(self._running and self._connection is not None)

        # Dynamically adding a method to aiosqlite.Connection at runtime.
        # Type checkers can't understand this monkey-patch, so we suppress the
        # "attr-defined" error that would otherwise be raised.
        _aiosqlite.Connection.is_alive = _is_alive  # type: ignore[attr-defined]

    _aiosqlite_patched = True


@asynccontextmanager
async def _connect(
    *, db_path: Path | None = None,
) -> AsyncIterator[aiosqlite.Connection]:
    """Import aiosqlite, apply the compatibility patch, and connect.

    Centralizes the deferred import + patch + connect sequence used by every
    database function in this module.

    Args:
        db_path: Explicit path to the SQLite database.  When ``None``
            (the default) falls back to :func:`get_db_path` — preserving
            existing desktop / CLI behaviour.

    Yields:
        An open aiosqlite connection to the sessions database.
    """
    import aiosqlite as _aiosqlite

    _patch_aiosqlite()

    target = db_path or get_db_path()
    async with _aiosqlite.connect(str(target), timeout=30.0) as conn:
        yield conn


@asynccontextmanager
async def get_checkpointer(
    *, db_path: Path | None = None,
) -> AsyncIterator[AsyncSqliteSaver]:
    """Get AsyncSqliteSaver for the global sessions database.

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

    Yields:
        AsyncSqliteSaver instance for checkpoint persistence.
    """
    from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver

    _patch_aiosqlite()

    target = db_path or get_db_path()
    async with AsyncSqliteSaver.from_conn_string(str(target)) as checkpointer:
        yield checkpointer


class ThreadInfo(TypedDict):
    """Thread metadata returned by `list_threads`."""

    thread_id: str
    """Unique identifier for the thread."""

    agent_name: str | None
    """Name of the agent that owns the thread."""

    updated_at: str | None
    """ISO timestamp of the last update."""

    title: NotRequired[str]
    """Display title — custom title if set, otherwise first human message truncated to ~50 chars."""

    message_count: NotRequired[int]
    """Number of messages in the thread."""

    status: NotRequired[str]
    """Live thread status from the state machine (e.g. 'idle', 'streaming')."""


def format_timestamp(iso_timestamp: str | None) -> str:
    """Format ISO timestamp for display (e.g., 'Dec 30, 6:10pm').

    Args:
        iso_timestamp: ISO 8601 timestamp string, or `None`.

    Returns:
        Formatted timestamp string or empty string if invalid.
    """
    if not iso_timestamp:
        return ""
    try:
        dt = datetime.fromisoformat(iso_timestamp).astimezone()
        return (
            dt.strftime("%b %d, %-I:%M%p")
            .lower()
            .replace("am", "am")
            .replace("pm", "pm")
        )
    except (ValueError, TypeError):
        logger.debug(
            "Failed to parse timestamp %r; displaying as blank",
            iso_timestamp,
            exc_info=True,
        )
        return ""


def get_db_path() -> Path:
    """Get path to global database.

    Returns:
        Path to the SQLite database file.
    """
    db_dir = Path.home() / ".icarus"
    db_dir.mkdir(parents=True, exist_ok=True)
    return db_dir / "sessions.db"


async def _table_exists(conn: aiosqlite.Connection, table: str) -> bool:
    """Check if a table exists in the database.

    Returns:
        True if table exists, False otherwise.
    """
    query = "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?"
    async with conn.execute(query, (table,)) as cursor:
        return await cursor.fetchone() is not None


async def list_threads(
    agent_name: str | None = None,
    limit: int = 50,
    cursor: tuple[str, str] | None = None,
    include_message_count: bool = False,
    *,
    db_path: Path | None = None,
) -> list[ThreadInfo]:
    """List threads from checkpoints table with cursor-based pagination.

    Args:
        agent_name: Optional filter by agent name.
        limit: Maximum number of threads to return per page.
        cursor: Composite pagination cursor — a ``(updated_at, thread_id)``
            tuple from the last thread of the previous page.  Pass ``None``
            for the first page.
        include_message_count: Whether to include message counts.

    Returns:
        List of `ThreadInfo` dicts with `thread_id`, `agent_name`,
            `updated_at`, and optionally `message_count` and `status`.
    """
    async with _connect(db_path=db_path) as conn:
        # Return empty if table doesn't exist yet (fresh install)
        if not await _table_exists(conn, "checkpoints"):
            return []

        where = ""
        having = ""
        params_list: list[str | int] = []

        if agent_name:
            where = "WHERE json_extract(metadata, '$.agent_name') = ?"
            params_list.append(agent_name)

        if cursor:
            having = "HAVING (updated_at < ? OR (updated_at = ? AND thread_id < ?))"
            params_list.extend([cursor[0], cursor[0], cursor[1]])

        params_list.append(limit)

        query = f"""
            SELECT thread_id,
                   json_extract(metadata, '$.agent_name') as agent_name,
                   MAX(json_extract(metadata, '$.updated_at')) as updated_at
            FROM checkpoints
            {where}
            GROUP BY thread_id
            {having}
            ORDER BY updated_at DESC, thread_id DESC
            LIMIT ?
        """
        params = tuple(params_list)

        async with conn.execute(query, params) as db_cursor:
            rows = await db_cursor.fetchall()
            threads: list[ThreadInfo] = [
                ThreadInfo(thread_id=r[0], agent_name=r[1], updated_at=r[2])
                for r in rows
            ]

        # Fetch message counts and titles from checkpoints
        if include_message_count and threads:
            from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer

            serde = JsonPlusSerializer()
            custom_titles = await _get_custom_titles(
                conn, [t["thread_id"] for t in threads],
            )
            for thread in threads:
                count, title = await _extract_thread_metadata(
                    conn, thread["thread_id"], serde
                )
                thread["message_count"] = count
                # Custom title takes priority over derived first-message title
                custom = custom_titles.get(thread["thread_id"])
                if custom:
                    thread["title"] = custom
                elif title:
                    thread["title"] = title

        return threads


_MAX_TITLE_LENGTH = 50


async def _extract_thread_metadata(
    conn: aiosqlite.Connection,
    thread_id: str,
    serde: JsonPlusSerializer,
) -> tuple[int, str | None]:
    """Extract message count and title from the most recent checkpoint.

    Deserializes the checkpoint blob to count messages and find the first
    human message for use as a display title.

    Args:
        conn: Database connection.
        thread_id: The thread ID to inspect.
        serde: Serializer for decoding checkpoint data.

    Returns:
        Tuple of (message_count, title).  Title is the first human
        message truncated to ``_MAX_TITLE_LENGTH`` chars, or ``None``.
    """
    query = """
        SELECT type, checkpoint
        FROM checkpoints
        WHERE thread_id = ?
        ORDER BY checkpoint_id DESC
        LIMIT 1
    """
    async with conn.execute(query, (thread_id,)) as cursor:
        row = await cursor.fetchone()
        if not row or not row[0] or not row[1]:
            return 0, None

        type_str, checkpoint_blob = row
        try:
            data = serde.loads_typed((type_str, checkpoint_blob))
            channel_values = data.get("channel_values", {})
            messages = channel_values.get("messages", [])

            title = None
            for msg in messages:
                if getattr(msg, "type", None) == "human":
                    content = getattr(msg, "content", "")
                    if isinstance(content, str) and content.strip():
                        text = content.strip()
                        title = text[:_MAX_TITLE_LENGTH]
                        if len(text) > _MAX_TITLE_LENGTH:
                            title += "..."
                        break

            return len(messages), title
        except (ValueError, TypeError, KeyError):
            logger.warning(
                "Failed to deserialize checkpoint for thread %s; "
                "message count will show as 0",
                thread_id,
                exc_info=True,
            )
            return 0, None


async def _ensure_metadata_table(conn: aiosqlite.Connection) -> None:
    """Create the thread_metadata table if it doesn't exist."""
    await conn.execute("""
        CREATE TABLE IF NOT EXISTS thread_metadata (
            thread_id TEXT PRIMARY KEY,
            custom_title TEXT
        )
    """)


async def rename_thread(thread_id: str, title: str, *, db_path: Path | None = None) -> None:
    """Set a custom title for a thread.

    Args:
        thread_id: The thread to rename.
        title: The new display title.
        db_path: Explicit path to the SQLite database.  When ``None``
            (the default) falls back to :func:`get_db_path`.
    """
    async with _connect(db_path=db_path) as conn:
        await _ensure_metadata_table(conn)
        await conn.execute(
            "INSERT INTO thread_metadata (thread_id, custom_title) VALUES (?, ?) "
            "ON CONFLICT(thread_id) DO UPDATE SET custom_title = ?",
            (thread_id, title, title),
        )
        await conn.commit()


async def has_custom_title(thread_id: str, *, db_path: Path | None = None) -> bool:
    """Check if a thread already has a custom title.

    Used to guard auto-generation so it doesn't overwrite manual renames.

    Args:
        thread_id: The thread to check.
        db_path: Explicit path to the SQLite database.  When ``None``
            (the default) falls back to :func:`get_db_path`.

    Returns:
        True if a custom title exists for this thread.
    """
    async with _connect(db_path=db_path) as conn:
        if not await _table_exists(conn, "thread_metadata"):
            return False
        query = "SELECT 1 FROM thread_metadata WHERE thread_id = ? AND custom_title IS NOT NULL"
        async with conn.execute(query, (thread_id,)) as db_cursor:
            return await db_cursor.fetchone() is not None


async def _get_custom_titles(
    conn: aiosqlite.Connection, thread_ids: list[str],
) -> dict[str, str]:
    """Fetch custom titles for a batch of threads.

    Returns:
        Mapping of thread_id → custom_title for threads that have one.
    """
    if not thread_ids:
        return {}
    if not await _table_exists(conn, "thread_metadata"):
        return {}
    placeholders = ",".join("?" for _ in thread_ids)
    query = f"SELECT thread_id, custom_title FROM thread_metadata WHERE thread_id IN ({placeholders})"  # noqa: S608
    async with conn.execute(query, thread_ids) as db_cursor:
        return {row[0]: row[1] for row in await db_cursor.fetchall() if row[1]}


async def delete_thread(thread_id: str, *, db_path: Path | None = None) -> None:
    """Delete all checkpoint data for a thread.

    Removes rows from ``checkpoints``, ``checkpoint_writes``, and
    ``checkpoint_blobs`` so the thread no longer appears in listings
    and its history is permanently erased.

    Args:
        thread_id: The thread to delete.
        db_path: Explicit path to the SQLite database.  When ``None``
            (the default) falls back to :func:`get_db_path`.
    """
    async with _connect(db_path=db_path) as conn:
        for table in ("checkpoints", "checkpoint_writes", "checkpoint_blobs", "thread_metadata"):
            if await _table_exists(conn, table):
                await conn.execute(
                    f"DELETE FROM {table} WHERE thread_id = ?",  # noqa: S608
                    (thread_id,),
                )
        await conn.commit()


_DEFAULT_THREAD_LIMIT = 50


def get_thread_limit() -> int:
    """Read the thread listing limit from `DA_CLI_RECENT_THREADS`.

    Falls back to `_DEFAULT_THREAD_LIMIT` when the variable is unset or contains
    a non-integer value. The result is clamped to a minimum of 1.

    Returns:
        Number of threads to display.
    """
    import os

    raw = os.environ.get("DA_CLI_RECENT_THREADS")
    if raw is None:
        return _DEFAULT_THREAD_LIMIT
    try:
        return max(1, int(raw))
    except ValueError:
        logger.warning(
            "Invalid DA_CLI_RECENT_THREADS value %r, using default %d",
            raw,
            _DEFAULT_THREAD_LIMIT,
        )
        return _DEFAULT_THREAD_LIMIT


