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

from __future__ import annotations

import asyncio
import contextlib
import logging
import sqlite3
from contextlib import asynccontextmanager
from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING, Any, NamedTuple, NotRequired, TypedDict, cast

from deepagents_code._constants import SYSTEM_MESSAGE_PREFIX

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

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

    from deepagents_code.output import OutputFormat

logger = logging.getLogger(__name__)

_aiosqlite_patched = False
_jsonplus_serializer: JsonPlusSerializer | None = None
_message_count_cache: dict[str, tuple[str | None, int]] = {}
_MAX_MESSAGE_COUNT_CACHE = 4096
_initial_prompt_cache: dict[str, tuple[str | None, str | None]] = {}
_MAX_INITIAL_PROMPT_CACHE = 4096
_recent_threads_cache: dict[tuple[str | None, int], list[ThreadInfo]] = {}
_MAX_RECENT_THREADS_CACHE_KEYS = 16


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  # ty: ignore[unresolved-attribute]

    _aiosqlite_patched = True


async def _drain_aiosqlite_worker(conn: aiosqlite.Connection) -> None:
    """Join the aiosqlite worker thread after its connection is closed.

    `aiosqlite.Connection` wraps a daemon `Thread` (`conn._thread`) that
    drains its tx queue independently of the caller's event loop. The
    library's `close()` puts a stop sentinel on the queue and awaits the
    sentinel's future, but does not explicitly join the worker thread.

    If the connection is leaked (no explicit close) and the surrounding
    event loop has already shut down, the worker can still pop a queued
    item (typically from `Connection.__del__` calling `stop()`) and call
    `future.get_loop().call_soon_threadsafe(...)` on the closed loop. That
    raises `RuntimeError: Event loop is closed`, which pytest surfaces as
    `PytestUnhandledThreadExceptionWarning` (and GitHub Actions then
    surfaces as a workflow annotation).

    Explicitly joining the worker thread after close guarantees it has
    exited before this coroutine returns, eliminating the race for any
    connection routed through `_connect` / `get_checkpointer`.
    """
    worker = getattr(conn, "_thread", None)
    if worker is None or not worker.is_alive():
        return
    # `RuntimeError` covers the "thread was never started" case; treat as
    # already drained.
    with contextlib.suppress(RuntimeError):
        await asyncio.to_thread(worker.join, 5.0)


@asynccontextmanager
async def _connect() -> 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.

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

    _patch_aiosqlite()

    conn: aiosqlite.Connection | None = None
    try:
        async with _aiosqlite.connect(str(get_db_path()), timeout=30.0) as opened:
            conn = opened
            yield opened
    finally:
        if conn is not None:
            await _drain_aiosqlite_worker(conn)


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."""

    created_at: NotRequired[str | None]
    """ISO timestamp of thread creation (earliest checkpoint)."""

    git_branch: NotRequired[str | None]
    """Git branch active when the thread was created."""

    initial_prompt: NotRequired[str | None]
    """First human message in the thread."""

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

    latest_checkpoint_id: NotRequired[str | None]
    """Most recent checkpoint ID for cache invalidation."""

    cwd: NotRequired[str | None]
    """Working directory where the thread was last used."""


class _CheckpointSummary(NamedTuple):
    """Structured data extracted from a thread's latest checkpoint."""

    message_count: int | None
    """Number of messages inlined in the latest checkpoint, or `None`.

    `None` means the latest checkpoint did not inline the `messages` channel
    value, so the count is unknown and must be reconstructed from the `writes`
    table. This happens when `messages` uses a `DeltaChannel` (a LangGraph
    channel the deepagents SDK applies to `messages` as of v0.6) and the latest
    checkpoint falls between periodic snapshots. An `int` (including `0`) is a
    trustworthy count.
    """

    initial_prompt: str | None
    """First human prompt recovered from the latest checkpoint."""


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 format_relative_timestamp(iso_timestamp: str | None) -> str:
    """Format ISO timestamp as relative time (e.g., '5m ago', '2h ago').

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

    Returns:
        Relative time string or empty string if invalid.
    """
    if not iso_timestamp:
        return ""
    try:
        dt = datetime.fromisoformat(iso_timestamp).astimezone()
    except (ValueError, TypeError):
        logger.debug(
            "Failed to parse timestamp %r; displaying as blank",
            iso_timestamp,
            exc_info=True,
        )
        return ""

    delta = datetime.now(tz=dt.tzinfo) - dt
    seconds = int(delta.total_seconds())
    if seconds < 0:
        return "just now"
    if seconds < 60:  # noqa: PLR2004
        return f"{seconds}s ago"
    minutes = seconds // 60
    if minutes < 60:  # noqa: PLR2004
        return f"{minutes}m ago"
    hours = minutes // 60
    if hours < 24:  # noqa: PLR2004
        return f"{hours}h ago"
    days = hours // 24
    if days < 30:  # noqa: PLR2004
        return f"{days}d ago"
    if days < 365:  # noqa: PLR2004
        months = days // 30
        return f"{months}mo ago"
    years = days // 365
    return f"{years}y ago"


def format_path(path: str | None) -> str:
    """Format a filesystem path for display.

    Paths under the user's home directory are shown relative to `~`.
    All other paths are returned as-is.

    Args:
        path: Absolute filesystem path, or `None`.

    Returns:
        Formatted path string, or empty string if path is falsy.
    """
    if not path:
        return ""
    try:
        home = str(Path.home())
        if path == home:
            return "~"
        prefix = home + "/"
        if path.startswith(prefix):
            return "~/" + path[len(prefix) :]
    except (RuntimeError, KeyError, OSError):
        logger.debug(
            "Could not resolve home directory for path formatting", exc_info=True
        )
        return path
    else:
        return path


_db_path: Path | None = None


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

    The result is cached after the first successful call to avoid repeated
    filesystem operations.

    Returns:
        Path to the SQLite database file.
    """
    global _db_path  # noqa: PLW0603  # Module-level cache requires global statement
    if _db_path is not None:
        return _db_path
    from deepagents_code.model_config import DEFAULT_STATE_DIR

    DEFAULT_STATE_DIR.mkdir(parents=True, exist_ok=True)
    _db_path = DEFAULT_STATE_DIR / "sessions.db"
    return _db_path


def generate_thread_id() -> str:
    """Generate a new thread ID as a full UUID7 string.

    Returns:
        UUID7 string (time-ordered for natural sort by creation time).
    """
    from uuid_utils import uuid7

    return str(uuid7())


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 = 20,
    include_message_count: bool = False,
    sort_by: str = "updated",
    branch: str | None = None,
    cwd: str | None = None,
) -> list[ThreadInfo]:
    """List threads from checkpoints table.

    Args:
        agent_name: Optional filter by agent name.
        limit: Maximum number of threads to return.
        include_message_count: Whether to include message counts.
        sort_by: Sort field — `"updated"` or `"created"`.
        branch: Optional filter by git branch name.
        cwd: Optional filter by working directory. Only threads whose stored
            `cwd` metadata equals this path are returned. Matching is an
            exact string comparison — no path normalization, symlink
            resolution, or prefix matching. Threads without a stored `cwd`
            (older rows) are excluded.

    Returns:
        List of `ThreadInfo` dicts with `thread_id`, `agent_name`,
            `updated_at`, `created_at`, `latest_checkpoint_id`, `git_branch`,
            `cwd`, and optionally `message_count`.

    Raises:
        ValueError: If `sort_by` is not `"updated"` or `"created"`.
    """
    async with _connect() as conn:
        if not await _table_exists(conn, "checkpoints"):
            return []

        if sort_by not in {"updated", "created"}:
            msg = f"Invalid sort_by {sort_by!r}; expected 'updated' or 'created'"
            raise ValueError(msg)
        order_col = "created_at" if sort_by == "created" else "updated_at"

        where_clauses: list[str] = []
        params_list: list[str | int] = []

        if agent_name:
            where_clauses.append("json_extract(metadata, '$.agent_name') = ?")
            params_list.append(agent_name)
        if branch:
            where_clauses.append("json_extract(metadata, '$.git_branch') = ?")
            params_list.append(branch)
        if cwd:
            where_clauses.append("json_extract(metadata, '$.cwd') = ?")
            params_list.append(cwd)

        where_sql = f"WHERE {' AND '.join(where_clauses)}" if where_clauses else ""

        query = f"""
            SELECT thread_id,
                   json_extract(metadata, '$.agent_name') as agent_name,
                   MAX(json_extract(metadata, '$.updated_at')) as updated_at,
                   MAX(checkpoint_id) as latest_checkpoint_id,
                   MIN(json_extract(metadata, '$.updated_at')) as created_at,
                   MAX(json_extract(metadata, '$.git_branch')) as git_branch,
                   MAX(json_extract(metadata, '$.cwd')) as cwd
            FROM checkpoints
            {where_sql}
            GROUP BY thread_id
            ORDER BY {order_col} DESC
            LIMIT ?
        """  # noqa: S608  # where_sql/order_col derived from controlled internal values; user values use ? placeholders
        params: tuple = (*params_list, limit)

        async with conn.execute(query, params) as cursor:
            rows = await cursor.fetchall()
            threads: list[ThreadInfo] = [
                ThreadInfo(
                    thread_id=r[0],
                    agent_name=r[1],
                    updated_at=r[2],
                    latest_checkpoint_id=r[3],
                    created_at=r[4],
                    git_branch=r[5],
                    cwd=r[6],
                )
                for r in rows
            ]

        # Fetch message counts if requested
        if include_message_count and threads:
            await _populate_message_counts(conn, threads)

        # Only cache unfiltered results so the thread selector modal
        # doesn't receive branch-/cwd-filtered or differently-sorted data.
        if sort_by == "updated" and branch is None and cwd is None:
            _cache_recent_threads(agent_name, limit, threads)
        return threads


async def populate_thread_checkpoint_details(
    threads: list[ThreadInfo],
    *,
    include_message_count: bool = True,
    include_initial_prompt: bool = True,
) -> list[ThreadInfo]:
    """Populate checkpoint-derived fields for an existing thread list.

    This is used by the `/threads` modal to enrich rows in one background pass,
    so the latest checkpoint is fetched and deserialized at most once per row.

    Args:
        threads: Thread rows to enrich in place.
        include_message_count: Whether to populate `message_count`.
        include_initial_prompt: Whether to populate `initial_prompt`.

    Returns:
        The same list object with missing checkpoint-derived fields populated.
    """
    if not threads or (not include_message_count and not include_initial_prompt):
        return threads

    async with _connect() as conn:
        await _populate_checkpoint_fields(
            conn,
            threads,
            include_message_count=include_message_count,
            include_initial_prompt=include_initial_prompt,
        )
    return threads


async def prewarm_thread_message_counts(limit: int | None = None) -> None:
    """Prewarm thread selector cache for faster `/threads` open.

    Fetches a bounded list of recent threads and populates checkpoint-derived
    fields for currently visible columns into the in-memory cache. Intended to
    run in a background worker during app startup.

    Args:
        limit: Maximum threads to prewarm. Uses `get_thread_limit()` when `None`.
    """
    thread_limit = limit if limit is not None else get_thread_limit()
    if thread_limit < 1:
        return

    try:
        from deepagents_code.model_config import load_thread_config

        cfg = load_thread_config()
        threads = await list_threads(limit=thread_limit, include_message_count=False)
        if threads:
            await populate_thread_checkpoint_details(
                threads,
                include_message_count=cfg.columns.get("messages", False),
                include_initial_prompt=cfg.columns.get("initial_prompt", False),
            )
        _cache_recent_threads(None, thread_limit, threads)
    except (OSError, sqlite3.Error):
        logger.debug("Could not prewarm thread selector cache", exc_info=True)
    except Exception:
        logger.warning(
            "Unexpected error while prewarming thread selector cache",
            exc_info=True,
        )


def get_cached_threads(
    agent_name: str | None = None,
    limit: int | None = None,
) -> list[ThreadInfo] | None:
    """Get cached recent threads, if available.

    Args:
        agent_name: Optional agent-name filter key.
        limit: Maximum rows requested. Uses `get_thread_limit()` when `None`.

    Returns:
        Copy of cached rows when available, otherwise `None`.
    """

    def _copy_with_cached_counts(rows: list[ThreadInfo]) -> list[ThreadInfo]:
        copied_rows = _copy_threads(rows)
        apply_cached_thread_message_counts(copied_rows)
        apply_cached_thread_initial_prompts(copied_rows)
        return copied_rows

    thread_limit = limit if limit is not None else get_thread_limit()
    if thread_limit < 1:
        return None

    exact = _recent_threads_cache.get((agent_name, thread_limit))
    if exact is not None:
        return _copy_with_cached_counts(exact)

    best_key: tuple[str | None, int] | None = None
    for key in _recent_threads_cache:
        cache_agent, cache_limit = key
        if cache_agent != agent_name or cache_limit < thread_limit:
            continue
        if best_key is None or cache_limit < best_key[1]:
            best_key = key

    if best_key is None:
        return None

    return _copy_with_cached_counts(_recent_threads_cache[best_key][:thread_limit])


def apply_cached_thread_message_counts(threads: list[ThreadInfo]) -> int:
    """Apply cached message counts onto thread rows when freshness matches.

    Args:
        threads: Thread rows to mutate in place.

    Returns:
        Number of rows that were populated from cache.
    """
    populated = 0
    for thread in threads:
        if "message_count" in thread:
            continue
        thread_id = thread["thread_id"]
        freshness = _thread_freshness(thread)
        cached = _message_count_cache.get(thread_id)
        if cached is None or cached[0] != freshness:
            continue
        thread["message_count"] = cached[1]
        populated += 1
    return populated


def apply_cached_thread_initial_prompts(threads: list[ThreadInfo]) -> int:
    """Apply cached initial prompts onto thread rows when freshness matches.

    Args:
        threads: Thread rows to mutate in place.

    Returns:
        Number of rows that were populated from cache.
    """
    populated = 0
    for thread in threads:
        if "initial_prompt" in thread:
            continue
        thread_id = thread["thread_id"]
        freshness = _thread_freshness(thread)
        cached = _initial_prompt_cache.get(thread_id)
        if cached is None or cached[0] != freshness:
            continue
        thread["initial_prompt"] = cached[1]
        populated += 1
    return populated


async def _populate_message_counts(
    conn: aiosqlite.Connection,
    threads: list[ThreadInfo],
) -> None:
    """Fill `message_count` on thread rows with cache-aware lookup."""
    await _populate_checkpoint_fields(
        conn,
        threads,
        include_message_count=True,
        include_initial_prompt=False,
    )


async def _get_jsonplus_serializer() -> JsonPlusSerializer:
    """Return a cached JsonPlus serializer, loading it off the UI loop."""
    global _jsonplus_serializer  # noqa: PLW0603  # Module-level cache requires global statement
    if _jsonplus_serializer is not None:
        return _jsonplus_serializer

    loop = asyncio.get_running_loop()
    _jsonplus_serializer = await loop.run_in_executor(None, _create_jsonplus_serializer)
    return _jsonplus_serializer


def _create_jsonplus_serializer() -> JsonPlusSerializer:
    """Import and create a JsonPlus serializer.

    Returns:
        A ready `JsonPlusSerializer` instance.
    """
    from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer

    return JsonPlusSerializer()


def _cache_message_count(thread_id: str, freshness: str | None, count: int) -> None:
    """Cache a thread's message count with a freshness token."""
    if len(_message_count_cache) >= _MAX_MESSAGE_COUNT_CACHE and (
        thread_id not in _message_count_cache
    ):
        oldest = next(iter(_message_count_cache))
        _message_count_cache.pop(oldest, None)
    _message_count_cache[thread_id] = (freshness, count)


def _cache_initial_prompt(
    thread_id: str,
    freshness: str | None,
    initial_prompt: str | None,
) -> None:
    """Cache a thread's initial prompt with a freshness token."""
    if len(_initial_prompt_cache) >= _MAX_INITIAL_PROMPT_CACHE and (
        thread_id not in _initial_prompt_cache
    ):
        oldest = next(iter(_initial_prompt_cache))
        _initial_prompt_cache.pop(oldest, None)
    _initial_prompt_cache[thread_id] = (freshness, initial_prompt)


def _thread_freshness(thread: ThreadInfo) -> str | None:
    """Return a cache freshness token for a thread row.

    The token is checkpoint-granular (`latest_checkpoint_id`). The
    writes-reconstructed `message_count` includes pending writes on the latest
    checkpoint, which in principle can change without a new checkpoint ID — so
    this token does not capture intra-checkpoint write churn. In practice that
    is benign: dcode only mutates `messages` through the agent graph, and every
    batch of message writes culminates in a new checkpoint (each superstep,
    `aupdate_state`, interrupt, and cancellation all advance
    `latest_checkpoint_id`). The only window where a cached count can lag is
    opening the `/threads` selector mid-superstep against an actively streaming
    thread; the selector does not live-refresh, so that count stays put until
    the modal is reopened (by then a new checkpoint exists and the cache
    refreshes). Making the key write-sensitive would require probing the
    `writes` table for every row on every `list_threads`, which is not worth it
    for a cosmetic count.
    """
    return thread.get("latest_checkpoint_id") or thread.get("updated_at")


def _cache_recent_threads(
    agent_name: str | None,
    limit: int,
    threads: list[ThreadInfo],
) -> None:
    """Store a copy of recent thread rows for fast selector startup."""
    key = (agent_name, max(1, limit))
    if len(_recent_threads_cache) >= _MAX_RECENT_THREADS_CACHE_KEYS and (
        key not in _recent_threads_cache
    ):
        _recent_threads_cache.clear()
    _recent_threads_cache[key] = _copy_threads(threads)


def _copy_threads(threads: list[ThreadInfo]) -> list[ThreadInfo]:
    """Return shallow-copied thread rows."""
    return [ThreadInfo(**thread) for thread in threads]


async def _populate_checkpoint_fields(
    conn: aiosqlite.Connection,
    threads: list[ThreadInfo],
    *,
    include_message_count: bool,
    include_initial_prompt: bool,
) -> None:
    """Populate checkpoint-derived thread fields with a batched latest-row pass."""
    serde = await _get_jsonplus_serializer()

    # Phase 1: apply cache hits, collect threads that need DB fetch.
    uncached: list[ThreadInfo] = []
    for thread in threads:
        thread_id = thread["thread_id"]
        freshness = _thread_freshness(thread)
        needs_count = False
        needs_prompt = False

        if include_message_count:
            cached = _message_count_cache.get(thread_id)
            if cached is not None and cached[0] == freshness:
                thread["message_count"] = cached[1]
            else:
                needs_count = True

        if include_initial_prompt and "initial_prompt" not in thread:
            cached_prompt = _initial_prompt_cache.get(thread_id)
            if cached_prompt is not None and cached_prompt[0] == freshness:
                thread["initial_prompt"] = cached_prompt[1]
            else:
                needs_prompt = True

        if needs_count or needs_prompt:
            uncached.append(thread)

    if not uncached:
        return

    # Phase 2: batch-fetch all uncached threads.
    uncached_ids = [t["thread_id"] for t in uncached]
    batch_results: dict[str, _CheckpointSummary] = {}
    if include_message_count or include_initial_prompt:
        batch_results = await _load_latest_checkpoint_summaries_batch(
            conn, uncached_ids, serde
        )
    # `initial_prompt` cannot be recovered from the latest checkpoint alone:
    # `after_model` middleware (e.g., `ResumeStateMiddleware`) writes partial
    # checkpoints whose `channel_values` omit `messages`. Read the very first
    # write to the `messages` channel from the `writes` table instead — that
    # row holds the user's original input.
    prompt_results: dict[str, str | None] = {}
    if include_initial_prompt:
        prompt_results = await _load_initial_prompts_from_writes_batch(
            conn, uncached_ids, serde
        )

    # Phase 3: apply inline results, deferring threads whose latest checkpoint
    # does not inline the `messages` channel value. When `messages` uses a
    # `DeltaChannel` (LangGraph channel applied by the deepagents SDK as of
    # v0.6) the full list is only snapshotted into `channel_values` periodically,
    # so the latest checkpoint usually omits it; the count must then be
    # reconstructed from the `writes` table.
    needs_writes_count: list[str] = []
    for thread in uncached:
        thread_id = thread["thread_id"]
        freshness = _thread_freshness(thread)

        if include_message_count and "message_count" not in thread:
            summary = batch_results.get(thread_id)
            if summary is not None and summary.message_count is not None:
                thread["message_count"] = summary.message_count
                _cache_message_count(thread_id, freshness, summary.message_count)
            else:
                needs_writes_count.append(thread_id)
        if include_initial_prompt and "initial_prompt" not in thread:
            if thread_id in prompt_results:
                prompt = prompt_results[thread_id]
            else:
                prompt = batch_results.get(
                    thread_id, _CheckpointSummary(None, None)
                ).initial_prompt
            thread["initial_prompt"] = prompt
            _cache_initial_prompt(thread_id, freshness, prompt)

    # Phase 4: reconstruct counts for delta-channel threads from the `writes`
    # table by replaying the `messages` writes through the canonical reducer.
    if needs_writes_count:
        writes_counts = await _load_message_counts_from_writes_batch(
            conn, needs_writes_count, serde
        )
        uncached_by_id = {t["thread_id"]: t for t in uncached}
        for thread_id in needs_writes_count:
            count = writes_counts.get(thread_id, 0)
            thread = uncached_by_id[thread_id]
            thread["message_count"] = count
            _cache_message_count(thread_id, _thread_freshness(thread), count)


_SQLITE_MAX_VARIABLE_NUMBER = 500
"""Max `?` placeholders per SQL query.

SQLite limits how many `?` parameters a single query can have (default 999,
lower on some builds). If a user accumulates hundreds of threads and the
`/threads` modal fetches them all at once, the `IN (?, ?, ...)` clause could
exceed that limit. We chunk to this size to stay safe.
"""


async def _load_latest_checkpoint_summaries_batch(
    conn: aiosqlite.Connection,
    thread_ids: list[str],
    serde: JsonPlusSerializer,
) -> dict[str, _CheckpointSummary]:
    """Batch-load the latest checkpoint summary for multiple threads.

    Uses a window function to fetch the latest checkpoint per thread, issuing
    one query per chunk for SQLite variable-limit safety.

    Args:
        conn: Database connection.
        thread_ids: Thread IDs to look up.
        serde: Serializer for decoding checkpoint blobs.

    Returns:
        Dict mapping thread IDs to their checkpoint summaries.
    """
    if not thread_ids:
        return {}

    results: dict[str, _CheckpointSummary] = {}

    for start in range(0, len(thread_ids), _SQLITE_MAX_VARIABLE_NUMBER):
        chunk = thread_ids[start : start + _SQLITE_MAX_VARIABLE_NUMBER]
        placeholders = ",".join("?" * len(chunk))
        query = f"""
            SELECT thread_id, type, checkpoint FROM (
                SELECT thread_id, type, checkpoint,
                       ROW_NUMBER() OVER (
                           PARTITION BY thread_id ORDER BY checkpoint_id DESC
                       ) AS rn
                FROM checkpoints
                WHERE thread_id IN ({placeholders})
            ) WHERE rn = 1
        """  # noqa: S608  # placeholders built from len(chunk); user values use ? params
        async with conn.execute(query, chunk) as cursor:
            rows = await cursor.fetchall()

        loop = asyncio.get_running_loop()
        for row in rows:
            tid, type_str, checkpoint_blob = row
            if not type_str or not checkpoint_blob:
                results[tid] = _CheckpointSummary(
                    message_count=None, initial_prompt=None
                )
                continue
            try:
                data = await loop.run_in_executor(
                    None, serde.loads_typed, (type_str, checkpoint_blob)
                )
                results[tid] = _summarize_checkpoint(data)
            except Exception:
                logger.warning(
                    "Failed to deserialize checkpoint for thread %s; "
                    "message count and initial prompt may be incomplete",
                    tid,
                    exc_info=True,
                )
                results[tid] = _CheckpointSummary(
                    message_count=None, initial_prompt=None
                )

    return results


async def _load_initial_prompts_from_writes_batch(
    conn: aiosqlite.Connection,
    thread_ids: list[str],
    serde: JsonPlusSerializer,
) -> dict[str, str | None]:
    """Batch-load initial prompts from the LangGraph `writes` table.

    For each thread, returns the first human/user message extracted from the
    earliest write to the `messages` channel (ordered by `checkpoint_id` ASC,
    then `idx` ASC).

    Args:
        conn: Database connection.
        thread_ids: Thread IDs to look up.
        serde: Serializer for decoding write blobs.

    Returns:
        Dict mapping thread IDs to their initial prompt text. Threads with no
        write to the `messages` channel are absent from the result; threads
        whose first such write decoded but contained no human/user entry map
        to `None`.
    """
    if not thread_ids:
        return {}

    results: dict[str, str | None] = {}
    loop = asyncio.get_running_loop()
    for start in range(0, len(thread_ids), _SQLITE_MAX_VARIABLE_NUMBER):
        chunk = thread_ids[start : start + _SQLITE_MAX_VARIABLE_NUMBER]
        placeholders = ",".join("?" * len(chunk))
        query = f"""
            SELECT thread_id, type, value FROM (
                SELECT thread_id, type, value,
                       ROW_NUMBER() OVER (
                           PARTITION BY thread_id
                           ORDER BY checkpoint_id ASC, idx ASC
                       ) AS rn
                FROM writes
                WHERE thread_id IN ({placeholders}) AND channel = 'messages'
            ) WHERE rn = 1
        """  # noqa: S608  # placeholders built from len(chunk); user values use ? params
        async with conn.execute(query, chunk) as cursor:
            rows = await cursor.fetchall()

        for row in rows:
            tid, type_str, value_blob = row
            if not type_str or not value_blob:
                continue
            try:
                messages = await loop.run_in_executor(
                    None, serde.loads_typed, (type_str, value_blob)
                )
            except Exception:
                logger.warning(
                    "Failed to deserialize initial messages write for thread %s",
                    tid,
                    exc_info=True,
                )
                continue
            if not isinstance(messages, list):
                continue
            results[tid] = _initial_prompt_from_messages(cast("list[object]", messages))

    return results


async def _load_message_counts_from_writes_batch(
    conn: aiosqlite.Connection,
    thread_ids: list[str],
    serde: JsonPlusSerializer,
) -> dict[str, int]:
    """Reconstruct message counts from the LangGraph `writes` table.

    For threads whose latest checkpoint does not inline the `messages` channel
    value — a `DeltaChannel` between snapshots, where the deepagents SDK (>= 0.6)
    applies LangGraph's `DeltaChannel` to `messages` — the full list is rebuilt
    by replaying every `messages` write, then counted. We replay through
    `add_messages` as a count-equivalent stand-in for the channel's actual
    reducer (`_messages_delta_reducer`): both dedup by ID and honor
    `RemoveMessage` / `REMOVE_ALL_MESSAGES`, so they produce the same final
    message set. The reducer is batching-invariant, so folding the write history
    from empty yields the same value LangGraph reconstructs from the latest
    snapshot plus subsequent deltas. An `Overwrite` write resets the accumulator
    to its value, matching the net effect of `DeltaChannel.replay_writes` (where
    the last `Overwrite` is the reset point).

    Only the root namespace (`checkpoint_ns = ''`) is counted, matching both the
    inline path and the conversation the `/threads` selector cares about;
    subgraph (subagent) writes under the same `thread_id` are excluded.

    Folding the *entire* write history (rather than walking the head
    checkpoint's parent chain) is intentional and matches what dcode shows when
    a thread is opened: it reads state via `aget_state` without a
    `checkpoint_id`, which applies pending writes (`apply_pending_writes=True`),
    so the latest checkpoint's not-yet-committed `messages` writes are part of
    the user-visible list and must be counted. dcode only ever appends to the
    latest checkpoint (no time travel, no `checkpoint_id`-targeted
    `aupdate_state`), so histories are linear and the full fold equals the
    head-of-chain reconstruction. A forked/abandoned branch (which dcode does
    not create) is the only case where this could over-count.

    Args:
        conn: Database connection.
        thread_ids: Thread IDs to look up.
        serde: Serializer for decoding write blobs.

    Returns:
        Dict mapping each thread ID with at least one decodable `messages`
            write to its reconstructed message count. Threads with no such
            writes are absent from the result.
    """
    if not thread_ids:
        return {}

    from langgraph.graph.message import add_messages
    from langgraph.types import Overwrite

    loop = asyncio.get_running_loop()
    # Accumulate one reduced message list per thread across chunks. Ordering by
    # (checkpoint_id, task_id, idx) replays deltas oldest-to-newest, matching
    # how LangGraph applies them on load.
    reduced: dict[str, list[Any]] = {}
    for start in range(0, len(thread_ids), _SQLITE_MAX_VARIABLE_NUMBER):
        chunk = thread_ids[start : start + _SQLITE_MAX_VARIABLE_NUMBER]
        placeholders = ",".join("?" * len(chunk))
        query = f"""
            SELECT thread_id, type, value
            FROM writes
            WHERE thread_id IN ({placeholders})
              AND checkpoint_ns = ''
              AND channel = 'messages'
            ORDER BY thread_id, checkpoint_id ASC, task_id ASC, idx ASC
        """  # noqa: S608  # placeholders built from len(chunk); user values use ? params
        async with conn.execute(query, chunk) as cursor:
            rows = await cursor.fetchall()

        for tid, type_str, value_blob in rows:
            if not type_str or not value_blob:
                continue
            try:
                delta = await loop.run_in_executor(
                    None, serde.loads_typed, (type_str, value_blob)
                )
                if isinstance(delta, Overwrite):
                    # Overwrite resets the channel; its value becomes the base.
                    value = delta.value
                    reduced[tid] = list(value) if isinstance(value, list) else []
                else:
                    # `add_messages` with a list left arg returns a list.
                    reduced[tid] = cast(
                        "list[Any]", add_messages(reduced.get(tid, []), delta)
                    )
            except Exception:
                logger.warning(
                    "Failed to replay messages write for thread %s; "
                    "message count may be inaccurate",
                    tid,
                    exc_info=True,
                )
                continue

    return {tid: len(messages) for tid, messages in reduced.items()}


def _summarize_checkpoint(data: object) -> _CheckpointSummary:
    """Extract message count and initial human prompt from checkpoint data.

    Returns:
        Structured summary for the decoded checkpoint payload.
    """
    messages = _checkpoint_messages(data)
    return _CheckpointSummary(
        message_count=len(messages) if messages is not None else None,
        initial_prompt=_initial_prompt_from_messages(messages or []),
    )


def _checkpoint_messages(data: object) -> list[object] | None:
    """Return inlined checkpoint messages, or `None` when not inlined.

    A `None` return distinguishes a checkpoint that omits the `messages`
    channel entirely (a `DeltaChannel` between snapshots, where the deepagents
    SDK applies LangGraph's `DeltaChannel` to `messages` as of v0.6) from one
    that inlines an empty list. The former requires reconstructing the count
    from the `writes` table; the latter is a genuine zero.
    """
    if not isinstance(data, dict):
        return None

    payload = cast("dict[str, object]", data)
    channel_values = payload.get("channel_values")
    if not isinstance(channel_values, dict):
        return None

    channel_values_dict = cast("dict[str, object]", channel_values)
    messages = channel_values_dict.get("messages")
    if not isinstance(messages, list):
        return None

    return cast("list[object]", messages)


def _initial_prompt_from_messages(messages: list[object]) -> str | None:
    """Return the first non-system human message content from a message list.

    Accepts both LangChain `HumanMessage` objects (with `type == "human"`) and
    plain dicts in OpenAI chat shape (`{"role": "user", "content": ...}`). The
    first write to the `messages` channel is the raw user input passed to the
    agent, which is preserved verbatim as a dict; subsequent writes are
    serialized `BaseMessage` instances produced after the model runs.

    Synthetic `[SYSTEM]`-prefixed human messages (e.g. an interrupt
    cancellation notice) are skipped so they never surface as a thread's prompt.
    """
    for msg in messages:
        if getattr(msg, "type", None) == "human":
            prompt = _coerce_prompt_text(getattr(msg, "content", None))
        elif isinstance(msg, dict):
            msg_dict = cast("dict[str, object]", msg)
            role = msg_dict.get("role")
            type_ = msg_dict.get("type")
            if role not in {"user", "human"} and type_ != "human":
                continue
            prompt = _coerce_prompt_text(msg_dict.get("content"))
        else:
            continue
        if prompt is not None and prompt.startswith(SYSTEM_MESSAGE_PREFIX):
            continue
        return prompt
    return None


def _coerce_prompt_text(content: object) -> str | None:
    """Normalize checkpoint message content into displayable text.

    Returns:
        Displayable prompt text, or `None` when the content is empty.
    """
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts: list[str] = []
        for part in content:
            if isinstance(part, dict):
                part_dict = cast("dict[str, object]", part)
                text = part_dict.get("text")
                parts.append(text if isinstance(text, str) else "")
            else:
                parts.append(str(part))
        joined = " ".join(parts).strip()
        return joined or None
    if content is None:
        return None
    return str(content)


async def get_most_recent(agent_name: str | None = None) -> str | None:
    """Get most recent thread_id, optionally filtered by agent.

    Returns:
        Most recent thread_id or None if no threads exist.
    """
    async with _connect() as conn:
        if not await _table_exists(conn, "checkpoints"):
            return None

        if agent_name:
            query = """
                SELECT thread_id FROM checkpoints
                WHERE json_extract(metadata, '$.agent_name') = ?
                ORDER BY checkpoint_id DESC
                LIMIT 1
            """
            params: tuple = (agent_name,)
        else:
            query = (
                "SELECT thread_id FROM checkpoints ORDER BY checkpoint_id DESC LIMIT 1"
            )
            params = ()

        async with conn.execute(query, params) as cursor:
            row = await cursor.fetchone()
            return row[0] if row else None


async def get_thread_agent(thread_id: str) -> str | None:
    """Get agent_name for a thread.

    Returns:
        Agent name associated with the thread, or None if not found.
    """
    async with _connect() as conn:
        if not await _table_exists(conn, "checkpoints"):
            return None

        query = """
            SELECT json_extract(metadata, '$.agent_name')
            FROM checkpoints
            WHERE thread_id = ?
            LIMIT 1
        """
        async with conn.execute(query, (thread_id,)) as cursor:
            row = await cursor.fetchone()
            return row[0] if row else None


async def get_thread_cwd(thread_id: str) -> str | None:
    """Get the most recently stored cwd for a thread.

    Args:
        thread_id: The thread whose stored cwd to look up.

    Returns:
        Most recent cwd for the thread, or None if not found.
    """
    async with _connect() as conn:
        if not await _table_exists(conn, "checkpoints"):
            return None

        query = """
            SELECT json_extract(metadata, '$.cwd')
            FROM checkpoints
            WHERE thread_id = ? AND json_extract(metadata, '$.cwd') IS NOT NULL
            ORDER BY checkpoint_id DESC
            LIMIT 1
        """
        async with conn.execute(query, (thread_id,)) as cursor:
            row = await cursor.fetchone()
            value = row[0] if row else None
            return value if isinstance(value, str) and value else None


async def thread_exists(thread_id: str) -> bool:
    """Check if a thread exists in checkpoints.

    Returns:
        True if thread exists, False otherwise.
    """
    async with _connect() as conn:
        if not await _table_exists(conn, "checkpoints"):
            return False

        query = "SELECT 1 FROM checkpoints WHERE thread_id = ? LIMIT 1"
        async with conn.execute(query, (thread_id,)) as cursor:
            row = await cursor.fetchone()
            return row is not None


async def find_similar_threads(thread_id: str, limit: int = 3) -> list[str]:
    """Find threads whose IDs start with the given prefix.

    Args:
        thread_id: Prefix to match against thread IDs.
        limit: Maximum number of matching threads to return.

    Returns:
        List of thread IDs that begin with the given prefix.
    """
    async with _connect() as conn:
        if not await _table_exists(conn, "checkpoints"):
            return []

        query = """
            SELECT DISTINCT thread_id
            FROM checkpoints
            WHERE thread_id LIKE ?
            ORDER BY thread_id
            LIMIT ?
        """
        prefix = thread_id + "%"
        async with conn.execute(query, (prefix, limit)) as cursor:
            rows = await cursor.fetchall()
            return [r[0] for r in rows]


async def delete_thread(thread_id: str) -> bool:
    """Delete thread checkpoints.

    Returns:
        True if thread was deleted, False if not found.
    """
    async with _connect() as conn:
        if not await _table_exists(conn, "checkpoints"):
            return False

        cursor = await conn.execute(
            "DELETE FROM checkpoints WHERE thread_id = ?", (thread_id,)
        )
        deleted = cursor.rowcount > 0
        if await _table_exists(conn, "writes"):
            await conn.execute("DELETE FROM writes WHERE thread_id = ?", (thread_id,))
        await conn.commit()
        if deleted:
            _message_count_cache.pop(thread_id, None)
            for key, rows in list(_recent_threads_cache.items()):
                filtered = [row for row in rows if row["thread_id"] != thread_id]
                _recent_threads_cache[key] = filtered
        return deleted


@asynccontextmanager
async def get_checkpointer() -> AsyncIterator[AsyncSqliteSaver]:
    """Get AsyncSqliteSaver for the global database.

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

    _patch_aiosqlite()

    saver: AsyncSqliteSaver | None = None
    try:
        async with AsyncSqliteSaver.from_conn_string(
            str(get_db_path())
        ) as checkpointer:
            saver = checkpointer
            yield checkpointer
    finally:
        if saver is not None:
            conn = getattr(saver, "conn", None)
            if conn is not None:
                await _drain_aiosqlite_worker(conn)


_DEFAULT_THREAD_LIMIT = 20


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


async def list_threads_command(
    agent_name: str | None = None,
    limit: int | None = None,
    sort_by: str | None = None,
    branch: str | None = None,
    cwd: str | None = None,
    verbose: bool = False,
    relative: bool | None = None,
    *,
    output_format: OutputFormat = "text",
) -> None:
    """CLI handler for `deepagents threads list`.

    Fetches and displays a table of recent conversation threads, optionally
    filtered by agent name, git branch, or working directory.

    Args:
        agent_name: Only show threads belonging to this agent.

            When `None`, threads for all agents are shown.
        limit: Maximum number of threads to display.

            When `None`, reads from `DA_CLI_RECENT_THREADS` or falls back to
            the default.
        sort_by: Sort field — `"updated"` or `"created"`.

            When `None`, reads from config (`~/.deepagents/config.toml`).
        branch: Only show threads from this git branch.
        cwd: Only show threads whose stored `cwd` metadata equals this path
            (exact string match — no normalization or prefix matching). When
            `None`, no cwd filter is applied. Threads without a stored `cwd`
            (older rows) are excluded when this is set.
        verbose: When `True`, show all columns (branch, created, prompt).
        relative: Show timestamps as relative time (e.g., '5m ago').

            When `None`, reads from config (`~/.deepagents/config.toml`).
        output_format: Output format — `'text'` (Rich) or `'json'`.
    """
    from deepagents_code.model_config import (
        load_thread_relative_time,
        load_thread_sort_order,
    )

    if sort_by is None:
        raw = load_thread_sort_order()
        sort_by = "created" if raw == "created_at" else "updated"
    if relative is None:
        relative = load_thread_relative_time()

    fmt_ts = format_relative_timestamp if relative else format_timestamp

    limit = get_thread_limit() if limit is None else max(1, limit)

    threads = await list_threads(
        agent_name,
        limit=limit,
        include_message_count=True,
        sort_by=sort_by,
        branch=branch,
        cwd=cwd,
    )

    if verbose and threads:
        await populate_thread_checkpoint_details(
            threads, include_message_count=False, include_initial_prompt=True
        )

    if output_format == "json":
        from deepagents_code.output import write_json

        write_json("threads list", list(threads))
        return

    from rich.markup import escape as escape_markup
    from rich.table import Table

    from deepagents_code import theme
    from deepagents_code.config import console

    if not threads:
        filters = []
        if agent_name:
            filters.append(f"agent '{escape_markup(agent_name)}'")
        if branch:
            filters.append(f"branch '{escape_markup(branch)}'")
        if cwd:
            filters.append(f"cwd '{escape_markup(cwd)}'")
        if filters:
            console.print(
                f"[yellow]No threads found for {' and '.join(filters)}.[/yellow]"
            )
        else:
            console.print("[yellow]No threads found.[/yellow]")
        if cwd:
            # Older threads predating cwd-metadata storage are silently
            # filtered out by the `json_extract` equality match. Tell the
            # user explicitly when unfiltered rows exist so they don't think
            # their history is gone.
            unfiltered = await list_threads(
                agent_name,
                limit=limit,
                include_message_count=False,
                sort_by=sort_by,
                branch=branch,
            )
            legacy_count = sum(1 for t in unfiltered if not t.get("cwd"))
            if legacy_count:
                console.print(
                    f"[dim]{legacy_count} older thread"
                    f"{'s have' if legacy_count != 1 else ' has'} "
                    "no cwd metadata; run without --cwd to see "
                    f"{'them' if legacy_count != 1 else 'it'}.[/dim]"
                )
        console.print("[dim]Start a conversation with: deepagents[/dim]")
        return

    title_parts = []
    if agent_name:
        title_parts.append(f"agent '{escape_markup(agent_name)}'")
    if branch:
        title_parts.append(f"branch '{escape_markup(branch)}'")
    if cwd:
        title_parts.append(f"cwd '{escape_markup(cwd)}'")

    title_filter = f" for {' and '.join(title_parts)}" if title_parts else ""
    sort_label = "created" if sort_by == "created" else "updated"
    title = f"Recent Threads{title_filter} (last {limit}, by {sort_label})"

    table = Table(title=title, show_header=True, header_style=f"bold {theme.PRIMARY}")
    table.add_column("Thread ID", style="bold")
    table.add_column("Agent")
    table.add_column("Messages", justify="right")
    if verbose:
        table.add_column("Created")
    table.add_column("Updated" if sort_by == "updated" else "Last Used")
    if verbose:
        table.add_column("Branch")
        table.add_column("Location")
        table.add_column("Prompt", max_width=40, no_wrap=True)

    prompt_max = 40

    for t in threads:
        row: list[str] = [
            t["thread_id"],
            t["agent_name"] or "unknown",
            str(t.get("message_count", 0)),
        ]
        if verbose:
            row.append(fmt_ts(t.get("created_at")))
        row.append(fmt_ts(t.get("updated_at")))
        if verbose:
            prompt = " ".join((t.get("initial_prompt") or "").split())
            if len(prompt) > prompt_max:
                prompt = prompt[: prompt_max - 3] + "..."
            row.extend(
                [
                    t.get("git_branch") or "",
                    format_path(t.get("cwd")),
                    prompt,
                ]
            )
        table.add_row(*row)

    console.print()
    console.print(table)
    if len(threads) >= limit:
        console.print(
            f"[dim]Showing last {limit} threads. "
            "Override with -n/--limit or DA_CLI_RECENT_THREADS.[/dim]"
        )
    console.print()


async def delete_thread_command(
    thread_id: str,
    *,
    dry_run: bool = False,
    output_format: OutputFormat = "text",
) -> None:
    """CLI handler for: deepagents threads delete.

    Args:
        thread_id: ID of the thread to delete.
        dry_run: If `True`, print what would happen without making changes.
        output_format: Output format — `'text'` (Rich) or `'json'`.
    """
    if dry_run:
        exists = await thread_exists(thread_id)
        if output_format == "json":
            from deepagents_code.output import write_json

            write_json(
                "threads delete",
                {"thread_id": thread_id, "exists": exists, "dry_run": True},
            )
            return

        from rich.markup import escape as escape_markup

        from deepagents_code.config import console

        escaped_id = escape_markup(thread_id)
        if exists:
            console.print(f"Would delete thread '{escaped_id}'.")
        else:
            console.print(f"Thread '{escaped_id}' not found. Nothing to delete.")
        console.print("No changes made.", style="dim")
        return

    deleted = await delete_thread(thread_id)

    if output_format == "json":
        from deepagents_code.output import write_json

        write_json("threads delete", {"thread_id": thread_id, "deleted": deleted})
        return

    from rich.markup import escape as escape_markup

    from deepagents_code import theme
    from deepagents_code.config import console

    escaped_id = escape_markup(thread_id)
    if deleted:
        console.print(f"[green]Thread '{escaped_id}' deleted.[/green]")
    else:
        console.print(
            f"Thread '{escaped_id}' not found or already deleted.",
            style=theme.MUTED,
        )
