#!/usr/bin/env python3
"""CLI-first local persistence and context recovery for agent-continuity."""

from __future__ import annotations

import argparse
import base64
import datetime as dt
import hashlib
import json
import math
import os
import shutil
import sqlite3
import sys
import tempfile
import uuid
from pathlib import Path
from typing import Any, Iterable


APP_NAME = "agent-continuity"
EXPORT_FORMAT = "agent-continuity-export"
DEFAULT_BRANCH_NAME = "main"


DEFAULT_CONFIG = {
    "retention": {
        "session_max_age_days": None,
        "processed_spool_max_age_days": 30,
        "max_backups": 5,
    },
    "backup": {
        "directory": "backups",
    },
}


def now_iso() -> str:
    return dt.datetime.now(dt.timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def new_id(prefix: str) -> str:
    return f"{prefix}_{uuid.uuid4().hex}"


def json_text(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True)


def hash_text(value: str) -> str:
    return hashlib.sha256(value.encode("utf-8")).hexdigest()


def estimate_tokens(value: Any) -> int:
    """Conservative cross-provider estimate; exact tokenizer counts are adapter-specific."""
    raw = value if isinstance(value, str) else json_text(value)
    return max(1, math.ceil(len(raw.encode("utf-8")) / 3.0))


def default_home() -> Path:
    configured = os.environ.get("AGENT_CONTINUITY_HOME")
    if configured:
        return Path(configured).expanduser()

    return Path.home() / ".agents" / "data" / "agent-continutiy"


SCHEMA_SQL = """
CREATE TABLE IF NOT EXISTS schema_meta (
    key TEXT PRIMARY KEY,
    value TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS sessions (
    session_id TEXT PRIMARY KEY,
    title TEXT NOT NULL,
    status TEXT NOT NULL DEFAULT 'active',
    active_branch_id TEXT NOT NULL,
    metadata_json TEXT NOT NULL DEFAULT '{}',
    created_at TEXT NOT NULL,
    updated_at TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS branches (
    branch_id TEXT PRIMARY KEY,
    session_id TEXT NOT NULL REFERENCES sessions(session_id) ON DELETE CASCADE,
    parent_branch_id TEXT REFERENCES branches(branch_id) ON DELETE SET NULL,
    name TEXT NOT NULL,
    status TEXT NOT NULL DEFAULT 'active' CHECK (status IN ('active', 'paused', 'archived')),
    metadata_json TEXT NOT NULL DEFAULT '{}',
    created_at TEXT NOT NULL,
    updated_at TEXT NOT NULL,
    UNIQUE(session_id, name)
);

CREATE TABLE IF NOT EXISTS runs (
    run_id TEXT PRIMARY KEY,
    session_id TEXT NOT NULL REFERENCES sessions(session_id) ON DELETE CASCADE,
    branch_id TEXT NOT NULL,
    agent_id TEXT NOT NULL,
    provider TEXT,
    model TEXT,
    request_id TEXT,
    status TEXT NOT NULL DEFAULT 'running',
    error_json TEXT,
    started_at TEXT NOT NULL,
    completed_at TEXT
);

CREATE TABLE IF NOT EXISTS messages (
    message_id TEXT PRIMARY KEY,
    session_id TEXT NOT NULL REFERENCES sessions(session_id) ON DELETE CASCADE,
    run_id TEXT REFERENCES runs(run_id) ON DELETE SET NULL,
    branch_id TEXT NOT NULL,
    sequence INTEGER NOT NULL,
    role TEXT NOT NULL CHECK (role IN ('user', 'assistant', 'tool')),
    status TEXT NOT NULL CHECK (status IN ('pending', 'streaming', 'complete', 'partial', 'failed')),
    content_json TEXT NOT NULL,
    text_content TEXT NOT NULL,
    provider_message_id TEXT,
    idempotency_key TEXT,
    sha256 TEXT NOT NULL,
    created_at TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS events (
    event_id TEXT PRIMARY KEY,
    session_id TEXT NOT NULL REFERENCES sessions(session_id) ON DELETE CASCADE,
    run_id TEXT REFERENCES runs(run_id) ON DELETE SET NULL,
    branch_id TEXT NOT NULL,
    sequence INTEGER NOT NULL,
    event_type TEXT NOT NULL,
    payload_json TEXT NOT NULL,
    idempotency_key TEXT,
    sha256 TEXT NOT NULL,
    created_at TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS checkpoints (
    checkpoint_id TEXT PRIMARY KEY,
    session_id TEXT NOT NULL REFERENCES sessions(session_id) ON DELETE CASCADE,
    branch_id TEXT NOT NULL,
    parent_event_id TEXT,
    summary_json TEXT NOT NULL,
    token_estimate INTEGER NOT NULL,
    created_by TEXT,
    created_at TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS native_states (
    state_id TEXT PRIMARY KEY,
    session_id TEXT NOT NULL REFERENCES sessions(session_id) ON DELETE CASCADE,
    provider TEXT NOT NULL,
    kind TEXT NOT NULL,
    payload_json TEXT NOT NULL,
    portable INTEGER NOT NULL DEFAULT 0,
    created_at TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS artifacts (
    artifact_id TEXT PRIMARY KEY,
    session_id TEXT NOT NULL REFERENCES sessions(session_id) ON DELETE CASCADE,
    run_id TEXT REFERENCES runs(run_id) ON DELETE SET NULL,
    branch_id TEXT NOT NULL,
    kind TEXT NOT NULL,
    path TEXT NOT NULL,
    sha256 TEXT,
    size_bytes INTEGER,
    metadata_json TEXT NOT NULL DEFAULT '{}',
    created_at TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS leases (
    session_id TEXT NOT NULL REFERENCES sessions(session_id) ON DELETE CASCADE,
    branch_id TEXT NOT NULL,
    agent_id TEXT NOT NULL,
    lease_until TEXT NOT NULL,
    PRIMARY KEY (session_id, branch_id)
);

CREATE INDEX IF NOT EXISTS idx_messages_session_branch_sequence
    ON messages(session_id, branch_id, sequence);
CREATE INDEX IF NOT EXISTS idx_events_session_branch_sequence
    ON events(session_id, branch_id, sequence);
CREATE INDEX IF NOT EXISTS idx_checkpoints_session_branch_created
    ON checkpoints(session_id, branch_id, created_at);
CREATE INDEX IF NOT EXISTS idx_branches_session_updated
    ON branches(session_id, updated_at);
CREATE INDEX IF NOT EXISTS idx_artifacts_session_branch_created
    ON artifacts(session_id, branch_id, created_at);
CREATE UNIQUE INDEX IF NOT EXISTS idx_messages_idempotency
    ON messages(session_id, idempotency_key)
    WHERE idempotency_key IS NOT NULL;
CREATE UNIQUE INDEX IF NOT EXISTS idx_events_idempotency
    ON events(session_id, idempotency_key)
    WHERE idempotency_key IS NOT NULL;
"""


class Store:
    def __init__(self, home: Path | None = None, db_path: Path | None = None) -> None:
        self.home = (home or default_home()).expanduser().resolve()
        self.db_path = (db_path or self.home / "agent-continuity.db").expanduser().resolve()
        self.spool_dir = self.home / "spool"
        self.config_path = self.home / "config.json"

    def prepare_dirs(self) -> None:
        self.home.mkdir(parents=True, exist_ok=True)
        self.spool_dir.mkdir(parents=True, exist_ok=True)
        try:
            self.home.chmod(0o700)
        except OSError:
            pass

    def connect(self) -> sqlite3.Connection:
        self.prepare_dirs()
        conn = sqlite3.connect(str(self.db_path), timeout=5.0)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA foreign_keys = ON")
        conn.execute("PRAGMA busy_timeout = 5000")
        conn.execute("PRAGMA journal_mode = WAL")
        conn.execute("PRAGMA synchronous = FULL")
        conn.executescript(SCHEMA_SQL)
        try:
            conn.execute(
                "CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts "
                "USING fts5(message_id UNINDEXED, session_id UNINDEXED, text_content)"
            )
            conn.execute(
                "INSERT INTO schema_meta(key, value) VALUES('fts_enabled', 'true') "
                "ON CONFLICT(key) DO UPDATE SET value = excluded.value"
            )
        except sqlite3.OperationalError:
            conn.execute(
                "INSERT INTO schema_meta(key, value) VALUES('fts_enabled', 'false') "
                "ON CONFLICT(key) DO UPDATE SET value = excluded.value"
            )
        conn.commit()
        try:
            self.db_path.chmod(0o600)
        except OSError:
            pass
        return conn

    def read_config(self) -> dict[str, Any]:
        config = json.loads(json.dumps(DEFAULT_CONFIG))
        if not self.config_path.exists():
            return config
        loaded = json.loads(self.config_path.read_text(encoding="utf-8"))
        if not isinstance(loaded, dict):
            raise ValueError("config.json must contain a JSON object")
        merge_config(config, loaded)
        return config

    def write_config(self, config: dict[str, Any]) -> None:
        self.prepare_dirs()
        encoded = json.dumps(config, ensure_ascii=False, indent=2) + "\n"
        temporary = self.config_path.with_suffix(".json.tmp")
        temporary.write_text(encoded, encoding="utf-8")
        try:
            temporary.chmod(0o600)
        except OSError:
            pass
        os.replace(temporary, self.config_path)
        try:
            directory_fd = os.open(self.home, os.O_RDONLY)
            try:
                os.fsync(directory_fd)
            finally:
                os.close(directory_fd)
        except OSError:
            pass

    def spool(self, record: dict[str, Any]) -> Path:
        self.prepare_dirs()
        path = self.spool_dir / f"{record['session_id']}.jsonl"
        line = json_text(record) + "\n"
        with path.open("a", encoding="utf-8") as handle:
            handle.write(line)
            handle.flush()
            os.fsync(handle.fileno())
        try:
            path.chmod(0o600)
        except OSError:
            pass
        return path


def merge_config(target: dict[str, Any], source: dict[str, Any]) -> None:
    for key, value in source.items():
        if isinstance(value, dict) and isinstance(target.get(key), dict):
            merge_config(target[key], value)
        else:
            target[key] = value


def parse_json_or_string(value: str) -> Any:
    try:
        return json.loads(value)
    except json.JSONDecodeError:
        return value


def config_command(store: Store, args: argparse.Namespace) -> dict[str, Any]:
    config = store.read_config()
    if args.action == "show":
        config["config_path"] = str(store.config_path)
        return config
    if args.action != "set":
        raise ValueError(f"unknown config action: {args.action}")
    parts = [part for part in args.key.split(".") if part]
    if not parts:
        raise ValueError("config key must not be empty")
    target: dict[str, Any] = config
    for part in parts[:-1]:
        if not isinstance(target.get(part), dict):
            target[part] = {}
        target = target[part]
    target[parts[-1]] = parse_json_or_string(args.value)
    store.write_config(config)
    return {"config_path": str(store.config_path), "key": args.key, "value": target[parts[-1]]}


def backup_database(store: Store, args: argparse.Namespace) -> dict[str, Any]:
    config = store.read_config()
    configured_directory = Path(config["backup"]["directory"]).expanduser()
    backup_directory = configured_directory if configured_directory.is_absolute() else store.home / configured_directory
    destination = Path(args.output).expanduser() if args.output else backup_directory / (
        f"agent-continuity-{dt.datetime.now(dt.timezone.utc).strftime('%Y%m%dT%H%M%SZ')}.db"
    )
    destination.parent.mkdir(parents=True, exist_ok=True)
    if destination.resolve() == store.db_path:
        raise ValueError("backup output must differ from the live database")
    if destination.exists() and not args.force:
        raise ValueError(f"backup output already exists: {destination}; use --force to replace it")
    with store.connect() as source:
        target = sqlite3.connect(str(destination))
        try:
            source.backup(target)
            target.commit()
        finally:
            target.close()
    try:
        destination.chmod(0o600)
    except OSError:
        pass

    keep = args.keep if args.keep is not None else int(config["retention"].get("max_backups", 5))
    backups = sorted(backup_directory.glob("agent-continuity-*.db"), key=lambda path: path.stat().st_mtime, reverse=True)
    removed = []
    for old in backups[max(keep, 0):]:
        if old.resolve() == destination.resolve():
            continue
        old.unlink()
        removed.append(str(old))
    return {"backup": str(destination), "kept": max(keep, 0), "removed": removed}


def parse_iso(value: str) -> dt.datetime:
    return dt.datetime.fromisoformat(value.replace("Z", "+00:00"))


def prune(store: Store, args: argparse.Namespace) -> dict[str, Any]:
    config = store.read_config()
    session_days = args.older_than_days
    if session_days is None:
        session_days = config["retention"].get("session_max_age_days")
    spool_days = args.processed_spool_days
    if spool_days is None:
        spool_days = config["retention"].get("processed_spool_max_age_days")
    now = dt.datetime.now(dt.timezone.utc)
    session_candidates: list[dict[str, Any]] = []
    if session_days is not None:
        cutoff = now - dt.timedelta(days=int(session_days))
        with store.connect() as conn:
            rows = conn.execute(
                "SELECT session_id, title, status, updated_at FROM sessions "
                "WHERE status IN ('completed', 'archived') ORDER BY updated_at ASC"
            ).fetchall()
            session_candidates = [
                dict(row)
                for row in rows
                if parse_iso(row["updated_at"]) < cutoff
            ]

    processed_candidates: list[Path] = []
    if spool_days is not None and store.spool_dir.joinpath("processed").exists():
        spool_cutoff = now - dt.timedelta(days=int(spool_days))
        for path in store.spool_dir.joinpath("processed").glob("*.jsonl"):
            if dt.datetime.fromtimestamp(path.stat().st_mtime, dt.timezone.utc) < spool_cutoff:
                processed_candidates.append(path)

    removed_sessions: list[str] = []
    removed_spool: list[str] = []
    if args.apply:
        with store.connect() as conn:
            for candidate in session_candidates:
                if has_fts(conn):
                    conn.execute("DELETE FROM messages_fts WHERE session_id = ?", (candidate["session_id"],))
                conn.execute("DELETE FROM sessions WHERE session_id = ?", (candidate["session_id"],))
                removed_sessions.append(candidate["session_id"])
        for path in processed_candidates:
            path.unlink()
            removed_spool.append(str(path))
    return {
        "dry_run": not args.apply,
        "session_candidates": session_candidates,
        "processed_spool_candidates": [str(path) for path in processed_candidates],
        "removed_sessions": removed_sessions,
        "removed_processed_spool": removed_spool,
    }


def row_json(row: sqlite3.Row | None) -> dict[str, Any] | None:
    return dict(row) if row is not None else None


def session_or_error(conn: sqlite3.Connection, session_id: str) -> sqlite3.Row:
    row = conn.execute("SELECT * FROM sessions WHERE session_id = ?", (session_id,)).fetchone()
    if row is None:
        raise ValueError(f"session not found: {session_id}")
    return row


def branch_or_error(conn: sqlite3.Connection, session_id: str, branch_id: str) -> sqlite3.Row:
    row = conn.execute(
        "SELECT * FROM branches WHERE session_id = ? AND branch_id = ?",
        (session_id, branch_id),
    ).fetchone()
    if row is None:
        raise ValueError(f"branch not found: {branch_id}")
    return row


def run_or_error(conn: sqlite3.Connection, run_id: str) -> sqlite3.Row:
    row = conn.execute("SELECT * FROM runs WHERE run_id = ?", (run_id,)).fetchone()
    if row is None:
        raise ValueError(f"run not found: {run_id}")
    return row


def next_sequence(conn: sqlite3.Connection, table: str, session_id: str, branch_id: str) -> int:
    if table not in {"messages", "events"}:
        raise ValueError("unsupported sequence table")
    row = conn.execute(
        f"SELECT COALESCE(MAX(sequence), 0) + 1 AS next_sequence FROM {table} "
        "WHERE session_id = ? AND branch_id = ?",
        (session_id, branch_id),
    ).fetchone()
    return int(row["next_sequence"])


def insert_event(
    conn: sqlite3.Connection,
    *,
    session_id: str,
    branch_id: str,
    event_type: str,
    payload: dict[str, Any],
    run_id: str | None = None,
    event_id: str | None = None,
    idempotency_key: str | None = None,
    created_at: str | None = None,
) -> dict[str, Any]:
    event_id = event_id or new_id("evt")
    created_at = created_at or now_iso()
    sequence = next_sequence(conn, "events", session_id, branch_id)
    payload_json = json_text(payload)
    digest = hash_text(payload_json)
    if idempotency_key:
        existing = conn.execute(
            "SELECT * FROM events WHERE session_id = ? AND idempotency_key = ?",
            (session_id, idempotency_key),
        ).fetchone()
        if existing is not None:
            return dict(existing)
    conn.execute(
        "INSERT INTO events(event_id, session_id, run_id, branch_id, sequence, event_type, "
        "payload_json, idempotency_key, sha256, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
        (
            event_id,
            session_id,
            run_id,
            branch_id,
            sequence,
            event_type,
            payload_json,
            idempotency_key,
            digest,
            created_at,
        ),
    )
    return {
        "event_id": event_id,
        "session_id": session_id,
        "run_id": run_id,
        "branch_id": branch_id,
        "sequence": sequence,
        "event_type": event_type,
        "payload_json": payload_json,
        "idempotency_key": idempotency_key,
        "sha256": digest,
        "created_at": created_at,
    }


def message_record(
    *,
    session_id: str,
    run_id: str,
    branch_id: str | None,
    role: str,
    status: str,
    text_content: str,
    provider_message_id: str | None,
    message_id: str | None,
    idempotency_key: str | None,
    created_at: str | None,
) -> dict[str, Any]:
    content = {"type": "text", "text": text_content}
    content_json = json_text(content)
    created_at = created_at or now_iso()
    message_id = message_id or new_id("msg")
    return {
        "event_id": new_id("evt"),
        "event_type": "message.recorded",
        "session_id": session_id,
        "run_id": run_id,
        "branch_id": branch_id,
        "message_id": message_id,
        "role": role,
        "status": status,
        "content_json": content_json,
        "text_content": text_content,
        "provider_message_id": provider_message_id,
        "idempotency_key": idempotency_key,
        "sha256": hash_text(content_json),
        "created_at": created_at,
    }


def insert_message_record(store: Store, record: dict[str, Any], allow_spool: bool = True) -> dict[str, Any]:
    try:
        with store.connect() as conn:
            session = session_or_error(conn, record["session_id"])
            branch_id = record.get("branch_id") or session["active_branch_id"]
            run_id = record.get("run_id")
            if run_id:
                run = conn.execute("SELECT * FROM runs WHERE run_id = ?", (run_id,)).fetchone()
                if run is None:
                    if allow_spool:
                        raise ValueError(f"run not found: {run_id}")
                    # Preserve the raw message even if the run record was lost.
                    run_id = None
                else:
                    if run["session_id"] != record["session_id"]:
                        raise ValueError("run and message belong to different sessions")
                    branch_id = run["branch_id"]

            if record.get("idempotency_key"):
                existing = conn.execute(
                    "SELECT * FROM messages WHERE session_id = ? AND idempotency_key = ?",
                    (record["session_id"], record["idempotency_key"]),
                ).fetchone()
                if existing is not None:
                    return {"status": "existing", "message": dict(existing)}

            existing = conn.execute(
                "SELECT * FROM messages WHERE message_id = ?", (record["message_id"],)
            ).fetchone()
            if existing is not None:
                return {"status": "existing", "message": dict(existing)}

            sequence = next_sequence(conn, "messages", record["session_id"], branch_id)
            conn.execute(
                "INSERT INTO messages(message_id, session_id, run_id, branch_id, sequence, role, status, "
                "content_json, text_content, provider_message_id, idempotency_key, sha256, created_at) "
                "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
                (
                    record["message_id"],
                    record["session_id"],
                    run_id,
                    branch_id,
                    sequence,
                    record["role"],
                    record["status"],
                    record["content_json"],
                    record["text_content"],
                    record.get("provider_message_id"),
                    record.get("idempotency_key"),
                    record["sha256"],
                    record.get("created_at") or now_iso(),
                ),
            )
            event = insert_event(
                conn,
                session_id=record["session_id"],
                branch_id=branch_id,
                run_id=run_id,
                event_type="message.recorded",
                payload={
                    "message_id": record["message_id"],
                    "role": record["role"],
                    "status": record["status"],
                },
                event_id=record.get("event_id"),
                idempotency_key=record.get("idempotency_key"),
                created_at=record.get("created_at"),
            )
            updated_at = now_iso()
            conn.execute(
                "UPDATE sessions SET updated_at = ? WHERE session_id = ?",
                (updated_at, record["session_id"]),
            )
            if has_fts(conn):
                conn.execute(
                    "INSERT INTO messages_fts(message_id, session_id, text_content) VALUES (?, ?, ?)",
                    (record["message_id"], record["session_id"], record["text_content"]),
                )
            return {
                "status": "recorded",
                "message_id": record["message_id"],
                "event_id": event["event_id"],
                "sequence": sequence,
                "session_id": record["session_id"],
                "branch_id": branch_id,
                "message_status": record["status"],
            }
    except ValueError as exc:
        if not allow_spool or "run not found" not in str(exc):
            raise
        spool_record = dict(record)
        spool_record["spool_error"] = str(exc)
        path = store.spool(spool_record)
        return {
            "status": "spooled",
            "session_id": record["session_id"],
            "message_id": record["message_id"],
            "spool_path": str(path),
        }
    except (sqlite3.Error, OSError) as exc:
        if not allow_spool:
            raise
        spool_record = dict(record)
        spool_record["spool_error"] = str(exc)
        path = store.spool(spool_record)
        return {
            "status": "spooled",
            "session_id": record["session_id"],
            "message_id": record["message_id"],
            "spool_path": str(path),
        }


def has_fts(conn: sqlite3.Connection) -> bool:
    row = conn.execute("SELECT value FROM schema_meta WHERE key = 'fts_enabled'").fetchone()
    return bool(row and row["value"] == "true")


def read_text(args: argparse.Namespace) -> str:
    provided = [args.text is not None, args.input_file is not None, args.input_value is not None]
    if sum(provided) != 1:
        raise ValueError("provide exactly one of --text, --input-file, or --input")
    if args.text is not None:
        return args.text
    if args.input_file is not None:
        return Path(args.input_file).expanduser().read_text(encoding="utf-8")
    if args.input_value == "-":
        return sys.stdin.read()
    return args.input_value


def add_text_input(parser: argparse.ArgumentParser) -> None:
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--text")
    group.add_argument("--input-file")
    group.add_argument("--input", dest="input_value")


def make_cursor(session_id: str, branch_id: str, before_sequence: int) -> str:
    raw = json_text(
        {"session_id": session_id, "branch_id": branch_id, "before_sequence": before_sequence}
    ).encode("utf-8")
    return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")


def parse_cursor(value: str) -> dict[str, Any]:
    padded = value + "=" * (-len(value) % 4)
    try:
        result = json.loads(base64.urlsafe_b64decode(padded).decode("utf-8"))
    except (ValueError, json.JSONDecodeError) as exc:
        raise ValueError("invalid cursor") from exc
    if not result.get("session_id"):
        raise ValueError("invalid cursor")
    return result


def message_dict(row: sqlite3.Row) -> dict[str, Any]:
    result = dict(row)
    result["content_json"] = json.loads(result["content_json"])
    return result


def checkpoint_dict(row: sqlite3.Row | None) -> dict[str, Any] | None:
    if row is None:
        return None
    result = dict(row)
    result["summary"] = json.loads(result.pop("summary_json"))
    return result


def latest_checkpoint(conn: sqlite3.Connection, session_id: str, branch_id: str) -> sqlite3.Row | None:
    return conn.execute(
        "SELECT * FROM checkpoints WHERE session_id = ? AND branch_id = ? "
        "ORDER BY created_at DESC LIMIT 1",
        (session_id, branch_id),
    ).fetchone()


def coverage(summary: dict[str, Any]) -> dict[str, bool]:
    keys = {
        "objective": ("objective", "goal"),
        "constraints": ("constraints",),
        "decisions": ("decisions",),
        "open_blockers": ("open_blockers", "blockers"),
        "next_action": ("next_action", "next_step"),
        "workspace": ("workspace",),
    }
    result: dict[str, bool] = {}
    for name, candidates in keys.items():
        result[name] = any(summary.get(key) not in (None, "", [], {}) for key in candidates)
    return result


def manifest(conn: sqlite3.Connection, session_id: str) -> dict[str, Any]:
    session = session_or_error(conn, session_id)
    branch_id = session["active_branch_id"]
    cp = latest_checkpoint(conn, session_id, branch_id)
    last_message = conn.execute(
        "SELECT * FROM messages WHERE session_id = ? AND branch_id = ? ORDER BY sequence DESC LIMIT 1",
        (session_id, branch_id),
    ).fetchone()
    event_count = conn.execute("SELECT COUNT(*) AS count FROM events WHERE session_id = ?", (session_id,)).fetchone()["count"]
    message_count = conn.execute("SELECT COUNT(*) AS count FROM messages WHERE session_id = ?", (session_id,)).fetchone()["count"]
    summary = json.loads(cp["summary_json"]) if cp else {}
    max_sequence_row = conn.execute(
        "SELECT COALESCE(MAX(sequence), 0) AS max_sequence FROM messages WHERE session_id = ? AND branch_id = ?",
        (session_id, branch_id),
    ).fetchone()
    max_sequence = int(max_sequence_row["max_sequence"])
    result = {
        "session_id": session_id,
        "title": session["title"],
        "status": session["status"],
        "active_branch_id": branch_id,
        "updated_at": session["updated_at"],
        "event_count": int(event_count),
        "message_count": int(message_count),
        "latest_checkpoint_id": cp["checkpoint_id"] if cp else None,
        "latest_message": {
            "role": last_message["role"],
            "status": last_message["status"],
            "sequence": last_message["sequence"],
            "created_at": last_message["created_at"],
        }
        if last_message
        else None,
        "objective": summary.get("objective", summary.get("goal")),
        "next_action": summary.get("next_action", summary.get("next_step")),
        "coverage": coverage(summary),
        "more_available": max_sequence > 0,
        "next_cursor": make_cursor(session_id, branch_id, max_sequence + 1),
    }
    return result


def emit(value: Any, pretty: bool = False) -> None:
    print(json.dumps(value, ensure_ascii=False, indent=2 if pretty else None, sort_keys=pretty))


def create_session(store: Store, args: argparse.Namespace) -> dict[str, Any]:
    session_id = args.session_id or new_id("sess")
    branch_id = new_id("br")
    title = args.title.strip()
    created_at = now_iso()
    with store.connect() as conn:
        conn.execute(
            "INSERT INTO sessions(session_id, title, status, active_branch_id, metadata_json, created_at, updated_at) "
            "VALUES (?, ?, 'active', ?, '{}', ?, ?)",
            (session_id, title, branch_id, created_at, created_at),
        )
        conn.execute(
            "INSERT INTO branches(branch_id, session_id, parent_branch_id, name, status, metadata_json, created_at, updated_at) "
            "VALUES (?, ?, NULL, ?, 'active', '{}', ?, ?)",
            (branch_id, session_id, DEFAULT_BRANCH_NAME, created_at, created_at),
        )
        insert_event(
            conn,
            session_id=session_id,
            branch_id=branch_id,
            event_type="session.created",
            payload={"title": title},
        )
    return {"session_id": session_id, "title": title, "status": "active", "branch_id": branch_id}


def begin_run(store: Store, args: argparse.Namespace) -> dict[str, Any]:
    run_id = new_id("run")
    started_at = now_iso()
    try:
        with store.connect() as conn:
            session = session_or_error(conn, args.session_id)
            branch_id = args.branch_id or session["active_branch_id"]
            branch_or_error(conn, args.session_id, branch_id)
            conn.execute(
                "INSERT INTO runs(run_id, session_id, branch_id, agent_id, provider, model, request_id, status, started_at) "
                "VALUES (?, ?, ?, ?, ?, ?, ?, 'running', ?)",
                (run_id, args.session_id, branch_id, args.agent, args.provider, args.model, args.request_id, started_at),
            )
            insert_event(
                conn,
                session_id=args.session_id,
                branch_id=branch_id,
                run_id=run_id,
                event_type="run.started",
                payload={"agent_id": args.agent, "provider": args.provider, "model": args.model},
            )
            lease = None
            if args.lease_ttl_seconds:
                lease = acquire_lease_conn(
                    conn,
                    session_id=args.session_id,
                    branch_id=branch_id,
                    agent_id=args.agent,
                    ttl_seconds=args.lease_ttl_seconds,
                )
                insert_event(
                    conn,
                    session_id=args.session_id,
                    branch_id=branch_id,
                    run_id=run_id,
                    event_type="lease.acquired",
                    payload={"agent_id": args.agent, "lease_until": lease["lease_until"]},
                )
        result = {"run_id": run_id, "session_id": args.session_id, "branch_id": branch_id, "status": "running"}
        if lease:
            result["lease"] = lease
        return result
    except (sqlite3.Error, OSError) as exc:
        branch_id = args.branch_id
        path = store.spool(
            {
                "kind": "run",
                "run_id": run_id,
                "session_id": args.session_id,
                "branch_id": branch_id,
                "agent_id": args.agent,
                "provider": args.provider,
                "model": args.model,
                "request_id": args.request_id,
                "status": "running",
                "started_at": started_at,
                "spool_error": str(exc),
            }
        )
        return {"status": "spooled", "run_id": run_id, "session_id": args.session_id, "spool_path": str(path)}


def finish_run(store: Store, args: argparse.Namespace) -> dict[str, Any]:
    completed_at = now_iso()
    error = json.loads(args.error_json) if args.error_json else None
    try:
        with store.connect() as conn:
            run = run_or_error(conn, args.run_id)
            conn.execute(
                "UPDATE runs SET status = ?, error_json = ?, completed_at = ? WHERE run_id = ?",
                (args.status, json_text(error) if error is not None else None, completed_at, args.run_id),
            )
            insert_event(
                conn,
                session_id=run["session_id"],
                branch_id=run["branch_id"],
                run_id=args.run_id,
                event_type="run.finished",
                payload={"status": args.status, "error": error},
            )
            if release_lease_conn(
                conn,
                session_id=run["session_id"],
                branch_id=run["branch_id"],
                agent_id=run["agent_id"],
            ):
                insert_event(
                    conn,
                    session_id=run["session_id"],
                    branch_id=run["branch_id"],
                    run_id=args.run_id,
                    event_type="lease.released",
                    payload={"agent_id": run["agent_id"], "reason": "run.finished"},
                )
    except (sqlite3.Error, OSError) as exc:
        if not args.session_id:
            raise ValueError("database unavailable; retry finish-run with --session-id") from exc
        path = store.spool(
            {
                "kind": "finish",
                "run_id": args.run_id,
                "session_id": args.session_id,
                "status": args.status,
                "error": error,
                "completed_at": completed_at,
                "spool_error": str(exc),
            }
        )
        return {"status": "spooled", "run_id": args.run_id, "spool_path": str(path)}
    return {"run_id": args.run_id, "status": args.status, "completed_at": completed_at}


def recover_finish_record(store: Store, record: dict[str, Any]) -> dict[str, Any]:
    with store.connect() as conn:
        run = run_or_error(conn, record["run_id"])
        completed_at = record.get("completed_at") or now_iso()
        conn.execute(
            "UPDATE runs SET status = ?, error_json = ?, completed_at = ? WHERE run_id = ?",
            (
                record["status"],
                json_text(record["error"]) if record.get("error") is not None else None,
                completed_at,
                record["run_id"],
            ),
        )
        insert_event(
            conn,
            session_id=run["session_id"],
            branch_id=run["branch_id"],
            run_id=record["run_id"],
            event_type="run.finished",
            payload={"status": record["status"], "error": record.get("error")},
        )
        if release_lease_conn(
            conn,
            session_id=run["session_id"],
            branch_id=run["branch_id"],
            agent_id=run["agent_id"],
        ):
            insert_event(
                conn,
                session_id=run["session_id"],
                branch_id=run["branch_id"],
                run_id=record["run_id"],
                event_type="lease.released",
                payload={"agent_id": run["agent_id"], "reason": "run.finished"},
            )
    return {"status": "recovered", "run_id": record["run_id"]}


def append_message(store: Store, args: argparse.Namespace, role: str) -> dict[str, Any]:
    try:
        with store.connect() as conn:
            run = run_or_error(conn, args.run_id)
        session_id = run["session_id"]
        branch_id = run["branch_id"]
    except (sqlite3.Error, OSError, ValueError):
        if not args.session_id:
            raise ValueError("run is unavailable; provide --session-id so the message can be spooled")
        session_id = args.session_id
        branch_id = args.branch_id
    text = read_text(args)
    record = message_record(
        session_id=session_id,
        run_id=args.run_id,
        branch_id=branch_id,
        role=role,
        status=args.status,
        text_content=text,
        provider_message_id=args.provider_message_id,
        message_id=args.message_id,
        idempotency_key=args.idempotency_key,
        created_at=None,
    )
    return insert_message_record(store, record)


def checkpoint(store: Store, args: argparse.Namespace) -> dict[str, Any]:
    if args.summary_json and args.summary_file:
        raise ValueError("use only one of --summary-json and --summary-file")
    if args.summary_file:
        summary = json.loads(Path(args.summary_file).expanduser().read_text(encoding="utf-8"))
    elif args.summary_json:
        summary = json.loads(args.summary_json)
    else:
        summary = {}
    if not isinstance(summary, dict):
        raise ValueError("checkpoint summary must be a JSON object")

    with store.connect() as conn:
        session = session_or_error(conn, args.session_id)
        branch_id = args.branch_id or session["active_branch_id"]
        branch_or_error(conn, args.session_id, branch_id)
        parent = conn.execute(
            "SELECT event_id FROM events WHERE session_id = ? AND branch_id = ? ORDER BY sequence DESC LIMIT 1",
            (args.session_id, branch_id),
        ).fetchone()
        if not summary:
            summary = auto_summary(
                conn,
                args.session_id,
                branch_id,
                max(args.source_budget_tokens, 1),
            )
        checkpoint_id = new_id("cp")
        created_at = now_iso()
        conn.execute(
            "INSERT INTO checkpoints(checkpoint_id, session_id, branch_id, parent_event_id, summary_json, "
            "token_estimate, created_by, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
            (
                checkpoint_id,
                args.session_id,
                branch_id,
                parent["event_id"] if parent else None,
                json_text(summary),
                estimate_tokens(summary),
                args.created_by,
                created_at,
            ),
        )
        event = insert_event(
            conn,
            session_id=args.session_id,
            branch_id=branch_id,
            event_type="checkpoint.created",
            payload={"checkpoint_id": checkpoint_id, "coverage": coverage(summary)},
        )
        conn.execute("UPDATE sessions SET updated_at = ? WHERE session_id = ?", (created_at, args.session_id))
    return {
        "checkpoint_id": checkpoint_id,
        "session_id": args.session_id,
        "branch_id": branch_id,
        "parent_event_id": parent["event_id"] if parent else None,
        "event_id": event["event_id"],
        "token_estimate": estimate_tokens(summary),
        "created_at": created_at,
    }


def create_branch(store: Store, args: argparse.Namespace) -> dict[str, Any]:
    branch_id = args.branch_id or new_id("br")
    created_at = now_iso()
    with store.connect() as conn:
        session = session_or_error(conn, args.session_id)
        source_branch_id = args.from_branch_id or session["active_branch_id"]
        branch_or_error(conn, args.session_id, source_branch_id)
        source_messages = conn.execute(
            "SELECT * FROM messages WHERE session_id = ? AND branch_id = ? ORDER BY sequence",
            (args.session_id, source_branch_id),
        ).fetchall()
        source_checkpoint = latest_checkpoint(conn, args.session_id, source_branch_id)
        conn.execute(
            "INSERT INTO branches(branch_id, session_id, parent_branch_id, name, status, metadata_json, created_at, updated_at) "
            "VALUES (?, ?, ?, ?, 'active', ?, ?, ?)",
            (
                branch_id,
                args.session_id,
                source_branch_id,
                args.name or branch_id,
                json_text(json.loads(args.metadata_json) if args.metadata_json else {}),
                created_at,
                created_at,
            ),
        )
        branch_event = insert_event(
            conn,
            session_id=args.session_id,
            branch_id=branch_id,
            event_type="branch.created",
            payload={"parent_branch_id": source_branch_id, "name": args.name or branch_id},
        )
        for source in source_messages:
            message_id = new_id("msg")
            sequence = next_sequence(conn, "messages", args.session_id, branch_id)
            conn.execute(
                "INSERT INTO messages(message_id, session_id, run_id, branch_id, sequence, role, status, "
                "content_json, text_content, provider_message_id, idempotency_key, sha256, created_at) "
                "VALUES (?, ?, NULL, ?, ?, ?, ?, ?, ?, ?, NULL, ?, ?)",
                (
                    message_id,
                    args.session_id,
                    branch_id,
                    sequence,
                    source["role"],
                    source["status"],
                    source["content_json"],
                    source["text_content"],
                    source["provider_message_id"],
                    source["sha256"],
                    source["created_at"],
                ),
            )
            insert_event(
                conn,
                session_id=args.session_id,
                branch_id=branch_id,
                event_type="message.forked",
                payload={"message_id": message_id, "source_message_id": source["message_id"]},
            )
            if has_fts(conn):
                conn.execute(
                    "INSERT INTO messages_fts(message_id, session_id, text_content) VALUES (?, ?, ?)",
                    (message_id, args.session_id, source["text_content"]),
                )
        if source_checkpoint is not None:
            checkpoint_id = new_id("cp")
            conn.execute(
                "INSERT INTO checkpoints(checkpoint_id, session_id, branch_id, parent_event_id, summary_json, "
                "token_estimate, created_by, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                (
                    checkpoint_id,
                    args.session_id,
                    branch_id,
                    branch_event["event_id"],
                    source_checkpoint["summary_json"],
                    source_checkpoint["token_estimate"],
                    source_checkpoint["created_by"],
                    created_at,
                ),
            )
            insert_event(
                conn,
                session_id=args.session_id,
                branch_id=branch_id,
                event_type="checkpoint.forked",
                payload={"checkpoint_id": checkpoint_id, "source_checkpoint_id": source_checkpoint["checkpoint_id"]},
            )
        if not args.no_switch:
            conn.execute(
                "UPDATE sessions SET active_branch_id = ?, updated_at = ? WHERE session_id = ?",
                (branch_id, created_at, args.session_id),
            )
    return {
        "branch_id": branch_id,
        "session_id": args.session_id,
        "parent_branch_id": source_branch_id,
        "switched": not args.no_switch,
        "copied_messages": len(source_messages),
    }


def switch_branch(store: Store, args: argparse.Namespace) -> dict[str, Any]:
    with store.connect() as conn:
        session_or_error(conn, args.session_id)
        branch = branch_or_error(conn, args.session_id, args.branch_id)
        timestamp = now_iso()
        conn.execute(
            "UPDATE sessions SET active_branch_id = ?, updated_at = ? WHERE session_id = ?",
            (args.branch_id, timestamp, args.session_id),
        )
        insert_event(
            conn,
            session_id=args.session_id,
            branch_id=args.branch_id,
            event_type="branch.switched",
            payload={"branch_id": args.branch_id},
        )
    return {"session_id": args.session_id, "branch_id": branch["branch_id"], "status": "active"}


def list_branches(store: Store, args: argparse.Namespace) -> list[dict[str, Any]]:
    with store.connect() as conn:
        session_or_error(conn, args.session_id)
        active_branch_id = conn.execute(
            "SELECT active_branch_id FROM sessions WHERE session_id = ?",
            (args.session_id,),
        ).fetchone()["active_branch_id"]
        rows = conn.execute(
            "SELECT * FROM branches WHERE session_id = ? ORDER BY created_at",
            (args.session_id,),
        ).fetchall()
        return [
            {
                "branch_id": row["branch_id"],
                "session_id": row["session_id"],
                "parent_branch_id": row["parent_branch_id"],
                "name": row["name"],
                "status": row["status"],
                "active": row["branch_id"] == active_branch_id,
                "created_at": row["created_at"],
                "updated_at": row["updated_at"],
            }
            for row in rows
        ]


def acquire_lease_conn(
    conn: sqlite3.Connection,
    *,
    session_id: str,
    branch_id: str,
    agent_id: str,
    ttl_seconds: int,
) -> dict[str, Any]:
    if ttl_seconds <= 0:
        raise ValueError("lease TTL must be positive")
    branch_or_error(conn, session_id, branch_id)
    now = dt.datetime.now(dt.timezone.utc)
    lease_until = (now + dt.timedelta(seconds=ttl_seconds)).replace(microsecond=0).isoformat().replace("+00:00", "Z")
    existing = conn.execute(
        "SELECT * FROM leases WHERE session_id = ? AND branch_id = ?",
        (session_id, branch_id),
    ).fetchone()
    if existing is not None and existing["lease_until"] > now_iso() and existing["agent_id"] != agent_id:
        raise ValueError(f"branch lease is held by {existing['agent_id']}")
    conn.execute(
        "INSERT INTO leases(session_id, branch_id, agent_id, lease_until) VALUES (?, ?, ?, ?) "
        "ON CONFLICT(session_id, branch_id) DO UPDATE SET agent_id = excluded.agent_id, lease_until = excluded.lease_until",
        (session_id, branch_id, agent_id, lease_until),
    )
    return {"session_id": session_id, "branch_id": branch_id, "agent_id": agent_id, "lease_until": lease_until}


def acquire_lease(store: Store, args: argparse.Namespace) -> dict[str, Any]:
    with store.connect() as conn:
        lease = acquire_lease_conn(
            conn,
            session_id=args.session_id,
            branch_id=args.branch_id,
            agent_id=args.agent,
            ttl_seconds=args.ttl_seconds,
        )
        insert_event(
            conn,
            session_id=args.session_id,
            branch_id=args.branch_id,
            event_type="lease.acquired",
            payload={"agent_id": args.agent, "lease_until": lease["lease_until"]},
        )
    return lease


def release_lease_conn(
    conn: sqlite3.Connection,
    *,
    session_id: str,
    branch_id: str,
    agent_id: str | None = None,
) -> bool:
    existing = conn.execute(
        "SELECT * FROM leases WHERE session_id = ? AND branch_id = ?",
        (session_id, branch_id),
    ).fetchone()
    if existing is None:
        return False
    if agent_id and existing["agent_id"] != agent_id:
        raise ValueError(f"branch lease is held by {existing['agent_id']}")
    conn.execute(
        "DELETE FROM leases WHERE session_id = ? AND branch_id = ?",
        (session_id, branch_id),
    )
    return True


def release_lease(store: Store, args: argparse.Namespace) -> dict[str, Any]:
    with store.connect() as conn:
        session_or_error(conn, args.session_id)
        branch_or_error(conn, args.session_id, args.branch_id)
        released = release_lease_conn(
            conn,
            session_id=args.session_id,
            branch_id=args.branch_id,
            agent_id=args.agent,
        )
        if released:
            insert_event(
                conn,
                session_id=args.session_id,
                branch_id=args.branch_id,
                event_type="lease.released",
                payload={"agent_id": args.agent},
            )
    return {"session_id": args.session_id, "branch_id": args.branch_id, "released": released}


def lease_status(store: Store, args: argparse.Namespace) -> dict[str, Any] | None:
    with store.connect() as conn:
        session_or_error(conn, args.session_id)
        branch_or_error(conn, args.session_id, args.branch_id)
        row = conn.execute(
            "SELECT * FROM leases WHERE session_id = ? AND branch_id = ?",
            (args.session_id, args.branch_id),
        ).fetchone()
        return dict(row) if row is not None else None


def file_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def add_artifact(store: Store, args: argparse.Namespace) -> dict[str, Any]:
    metadata = json.loads(args.metadata_json) if args.metadata_json else {}
    if not isinstance(metadata, dict):
        raise ValueError("artifact metadata must be a JSON object")
    artifact_path = Path(args.path).expanduser()
    sha256 = args.sha256
    size_bytes = args.size_bytes
    if artifact_path.is_file():
        sha256 = sha256 or file_sha256(artifact_path)
        size_bytes = size_bytes if size_bytes is not None else artifact_path.stat().st_size

    with store.connect() as conn:
        session = session_or_error(conn, args.session_id)
        branch_id = args.branch_id or session["active_branch_id"]
        branch_or_error(conn, args.session_id, branch_id)
        if args.run_id:
            run = run_or_error(conn, args.run_id)
            if run["session_id"] != args.session_id:
                raise ValueError("run and artifact belong to different sessions")
            branch_id = run["branch_id"]
        artifact_id = new_id("artifact")
        created_at = now_iso()
        conn.execute(
            "INSERT INTO artifacts(artifact_id, session_id, run_id, branch_id, kind, path, sha256, size_bytes, "
            "metadata_json, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                artifact_id,
                args.session_id,
                args.run_id,
                branch_id,
                args.kind,
                str(artifact_path),
                sha256,
                size_bytes,
                json_text(metadata),
                created_at,
            ),
        )
        insert_event(
            conn,
            session_id=args.session_id,
            branch_id=branch_id,
            run_id=args.run_id,
            event_type="artifact.added",
            payload={"artifact_id": artifact_id, "kind": args.kind, "path": str(artifact_path)},
        )
    return {
        "artifact_id": artifact_id,
        "session_id": args.session_id,
        "branch_id": branch_id,
        "path": str(artifact_path),
        "sha256": sha256,
        "size_bytes": size_bytes,
    }


def list_artifacts(store: Store, args: argparse.Namespace) -> list[dict[str, Any]]:
    with store.connect() as conn:
        session_or_error(conn, args.session_id)
        sql = "SELECT * FROM artifacts WHERE session_id = ?"
        params: list[Any] = [args.session_id]
        if args.branch_id:
            sql += " AND branch_id = ?"
            params.append(args.branch_id)
        sql += " ORDER BY created_at"
        rows = conn.execute(sql, params).fetchall()
        result = []
        for row in rows:
            item = dict(row)
            item["metadata"] = json.loads(item.pop("metadata_json"))
            result.append(item)
        return result


def select_messages(
    conn: sqlite3.Connection,
    *,
    session_id: str,
    branch_id: str,
    before_sequence: int | None,
    query: str | None,
    budget_tokens: int,
) -> list[dict[str, Any]]:
    sql = "SELECT * FROM messages WHERE session_id = ? AND branch_id = ?"
    params: list[Any] = [session_id, branch_id]
    if before_sequence is not None:
        sql += " AND sequence < ?"
        params.append(before_sequence)
    if query:
        terms = [term for term in query.split() if term]
        for term in terms:
            sql += " AND text_content LIKE ?"
            params.append(f"%{term}%")
    sql += " ORDER BY sequence DESC"
    rows = conn.execute(sql, params).fetchall()
    selected: list[dict[str, Any]] = []
    used = 0
    for row in rows:
        item = message_dict(row)
        cost = estimate_tokens(item)
        if selected and used + cost > budget_tokens:
            break
        if not selected and cost > budget_tokens:
            selected.append(item)
            break
        selected.append(item)
        used += cost
    selected.reverse()
    return selected


def truncate_text(value: str, limit: int) -> str:
    if len(value) <= limit:
        return value
    return value[: max(limit - 1, 1)] + "…"


def auto_summary(
    conn: sqlite3.Connection,
    session_id: str,
    branch_id: str,
    budget_tokens: int,
) -> dict[str, Any]:
    session = session_or_error(conn, session_id)
    previous = latest_checkpoint(conn, session_id, branch_id)
    summary: dict[str, Any] = {}
    if previous is not None:
        summary = json.loads(previous["summary_json"])

    messages = select_messages(
        conn,
        session_id=session_id,
        branch_id=branch_id,
        before_sequence=None,
        query=None,
        budget_tokens=budget_tokens,
    )
    latest_user = next((item for item in reversed(messages) if item["role"] == "user"), None)
    latest_assistant = next((item for item in reversed(messages) if item["role"] == "assistant"), None)
    summary["objective"] = summary.get("objective") or summary.get("goal") or session["title"]
    if latest_user:
        summary["latest_user_request"] = truncate_text(latest_user["text_content"], 800)
    if latest_assistant:
        summary["latest_assistant_result"] = truncate_text(latest_assistant["text_content"], 1200)
    summary["next_action"] = summary.get("next_action") or summary.get("next_step") or "继续处理最近一次用户请求"
    summary["source_message_ids"] = [item["message_id"] for item in messages]
    summary["auto_generated"] = True
    summary["source_budget_tokens"] = budget_tokens
    return summary


def resume(store: Store, args: argparse.Namespace) -> dict[str, Any]:
    with store.connect() as conn:
        session = session_or_error(conn, args.session_id)
        branch_id = session["active_branch_id"]
        cp = latest_checkpoint(conn, args.session_id, branch_id)
        base = manifest(conn, args.session_id)
        checkpoint_payload = checkpoint_dict(cp)
        base_payload = {"session": base, "checkpoint": checkpoint_payload, "messages": []}
        remaining = max(args.budget_tokens - estimate_tokens(base_payload), 1)
        selected = select_messages(
            conn,
            session_id=args.session_id,
            branch_id=branch_id,
            before_sequence=None,
            query=None,
            budget_tokens=remaining,
        )
        if selected:
            oldest = selected[0]["sequence"]
        else:
            oldest = int(base["latest_message"]["sequence"]) + 1 if base["latest_message"] else 1
        older = conn.execute(
            "SELECT 1 FROM messages WHERE session_id = ? AND branch_id = ? AND sequence < ? LIMIT 1",
            (args.session_id, branch_id, oldest),
        ).fetchone()
        return {
            "mode": "resume",
            "session": base,
            "checkpoint": checkpoint_payload,
            "messages": selected,
            "selected_tokens": estimate_tokens({"session": base, "checkpoint": checkpoint_payload, "messages": selected}),
            "more_available": older is not None,
            "next_cursor": make_cursor(args.session_id, branch_id, oldest),
        }


def expand(store: Store, args: argparse.Namespace) -> dict[str, Any]:
    cursor = parse_cursor(args.cursor)
    with store.connect() as conn:
        session = session_or_error(conn, cursor["session_id"])
        branch_id = cursor["branch_id"]
        if branch_id != session["active_branch_id"] and not args.allow_inactive_branch:
            raise ValueError("cursor points to an inactive branch; use --allow-inactive-branch to read it")
        before = int(cursor["before_sequence"])
        items = select_messages(
            conn,
            session_id=cursor["session_id"],
            branch_id=branch_id,
            before_sequence=before,
            query=args.query,
            budget_tokens=max(args.budget_tokens, 1),
        )
        oldest = items[0]["sequence"] if items else before
        older = conn.execute(
            "SELECT 1 FROM messages WHERE session_id = ? AND branch_id = ? AND sequence < ? LIMIT 1",
            (cursor["session_id"], branch_id, oldest),
        ).fetchone()
        return {
            "mode": "expand",
            "session_id": cursor["session_id"],
            "branch_id": branch_id,
            "query": args.query,
            "items": items,
            "selected_tokens": estimate_tokens(items),
            "more_available": older is not None,
            "next_cursor": make_cursor(cursor["session_id"], branch_id, oldest),
        }


def save_native_state(store: Store, args: argparse.Namespace) -> dict[str, Any]:
    payload = json.loads(args.payload_json)
    with store.connect() as conn:
        session = session_or_error(conn, args.session_id)
        state_id = new_id("state")
        created_at = now_iso()
        conn.execute(
            "INSERT INTO native_states(state_id, session_id, provider, kind, payload_json, portable, created_at) "
            "VALUES (?, ?, ?, ?, ?, ?, ?)",
            (state_id, args.session_id, args.provider, args.kind, json_text(payload), int(args.portable), created_at),
        )
        insert_event(
            conn,
            session_id=args.session_id,
            branch_id=session["active_branch_id"],
            event_type="native_state.saved",
            payload={"state_id": state_id, "provider": args.provider, "kind": args.kind, "portable": args.portable},
        )
    return {"state_id": state_id, "session_id": args.session_id, "provider": args.provider, "kind": args.kind}


def export_session(store: Store, args: argparse.Namespace) -> dict[str, Any]:
    with store.connect() as conn:
        session = session_or_error(conn, args.session_id)
        payload: dict[str, Any] = {
            "format": EXPORT_FORMAT,
            "session": {
                "session_id": session["session_id"],
                "title": session["title"],
                "status": session["status"],
                "active_branch_id": session["active_branch_id"],
                "metadata": json.loads(session["metadata_json"]),
            },
            "branches": [],
            "messages": [message_dict(row) for row in conn.execute("SELECT * FROM messages WHERE session_id = ? ORDER BY sequence", (args.session_id,))],
            "events": [],
            "checkpoints": [],
            "native_states": [],
            "artifacts": [],
        }
        for row in conn.execute("SELECT * FROM branches WHERE session_id = ? ORDER BY created_at", (args.session_id,)):
            item = dict(row)
            item["metadata"] = json.loads(item.pop("metadata_json"))
            payload["branches"].append(item)
        for row in conn.execute("SELECT * FROM events WHERE session_id = ? ORDER BY sequence", (args.session_id,)):
            item = dict(row)
            item["payload"] = json.loads(item.pop("payload_json"))
            payload["events"].append(item)
        for row in conn.execute("SELECT * FROM checkpoints WHERE session_id = ? ORDER BY created_at", (args.session_id,)):
            item = dict(row)
            item["summary"] = json.loads(item.pop("summary_json"))
            payload["checkpoints"].append(item)
        for row in conn.execute("SELECT * FROM native_states WHERE session_id = ? ORDER BY created_at", (args.session_id,)):
            item = dict(row)
            item["payload"] = json.loads(item.pop("payload_json"))
            item["portable"] = bool(item["portable"])
            payload["native_states"].append(item)
        for row in conn.execute("SELECT * FROM artifacts WHERE session_id = ? ORDER BY created_at", (args.session_id,)):
            item = dict(row)
            item["metadata"] = json.loads(item.pop("metadata_json"))
            payload["artifacts"].append(item)
    encoded = json.dumps(payload, ensure_ascii=False, indent=2)
    if args.output:
        Path(args.output).expanduser().write_text(encoded + "\n", encoding="utf-8")
        return {"session_id": args.session_id, "output": str(Path(args.output).expanduser()), "counts": {key: len(payload[key]) for key in ("messages", "events", "checkpoints")}}
    print(encoded)
    return {"session_id": args.session_id, "output": "stdout"}


def recover_run_record(store: Store, record: dict[str, Any]) -> dict[str, Any]:
    with store.connect() as conn:
        session = session_or_error(conn, record["session_id"])
        branch_id = record.get("branch_id") or session["active_branch_id"]
        branch_or_error(conn, record["session_id"], branch_id)
        existing = conn.execute("SELECT * FROM runs WHERE run_id = ?", (record["run_id"],)).fetchone()
        if existing is not None:
            return {"status": "existing", "run_id": record["run_id"]}
        conn.execute(
            "INSERT INTO runs(run_id, session_id, branch_id, agent_id, provider, model, request_id, status, started_at) "
            "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                record["run_id"],
                record["session_id"],
                branch_id,
                record["agent_id"],
                record.get("provider"),
                record.get("model"),
                record.get("request_id"),
                record.get("status", "running"),
                record.get("started_at") or now_iso(),
            ),
        )
        insert_event(
            conn,
            session_id=record["session_id"],
            branch_id=branch_id,
            run_id=record["run_id"],
            event_type="run.started",
            payload={"agent_id": record["agent_id"], "provider": record.get("provider"), "model": record.get("model")},
        )
    return {"status": "recovered", "run_id": record["run_id"]}


def recover_spool(store: Store) -> dict[str, Any]:
    processed = 0
    failed: list[str] = []
    processed_dir = store.spool_dir / "processed"
    processed_dir.mkdir(parents=True, exist_ok=True)
    for path in sorted(store.spool_dir.glob("*.jsonl")):
        try:
            records = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
            for record in sorted(records, key=lambda item: 0 if item.get("kind") == "run" else 1):
                if record.get("kind") == "run":
                    recover_run_record(store, record)
                elif record.get("kind") == "finish":
                    recover_finish_record(store, record)
                else:
                    insert_message_record(store, record, allow_spool=False)
            target = processed_dir / f"{path.stem}.{int(dt.datetime.now().timestamp())}.jsonl"
            shutil.move(str(path), str(target))
            processed += len(records)
        except (OSError, ValueError, sqlite3.Error, json.JSONDecodeError) as exc:
            failed.append(f"{path}: {exc}")
    return {"processed_records": processed, "failed_files": failed}


def verify(store: Store) -> dict[str, Any]:
    with store.connect() as conn:
        integrity = conn.execute("PRAGMA integrity_check").fetchone()[0]
        result = {
            "database": str(store.db_path),
            "integrity": integrity,
            "sessions": conn.execute("SELECT COUNT(*) FROM sessions").fetchone()[0],
            "branches": conn.execute("SELECT COUNT(*) FROM branches").fetchone()[0],
            "runs": conn.execute("SELECT COUNT(*) FROM runs").fetchone()[0],
            "messages": conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0],
            "events": conn.execute("SELECT COUNT(*) FROM events").fetchone()[0],
            "checkpoints": conn.execute("SELECT COUNT(*) FROM checkpoints").fetchone()[0],
            "artifacts": conn.execute("SELECT COUNT(*) FROM artifacts").fetchone()[0],
            "pending_runs": conn.execute("SELECT COUNT(*) FROM runs WHERE status = 'running'").fetchone()[0],
            "spool_files": len(list(store.spool_dir.glob("*.jsonl"))) if store.spool_dir.exists() else 0,
        }
    return result


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog=APP_NAME, description="CLI-first cross-agent task continuity")
    parser.add_argument("--home", type=Path, help="override the local data directory")
    parser.add_argument("--db", type=Path, help="override the SQLite database path")
    sub = parser.add_subparsers(dest="command", required=True)

    sub.add_parser("init", help="initialize the local database")

    create = sub.add_parser("create", help="create a task session")
    create.add_argument("--title", required=True)
    create.add_argument("--session-id")
    create.add_argument("--pretty", action="store_true")

    listing = sub.add_parser("list", help="list sessions")
    listing.add_argument("--pretty", action="store_true")

    begin = sub.add_parser("begin-run", help="start an Agent run")
    begin.add_argument("session_id")
    begin.add_argument("--agent", required=True)
    begin.add_argument("--provider")
    begin.add_argument("--model")
    begin.add_argument("--branch-id")
    begin.add_argument("--request-id")
    begin.add_argument("--lease-ttl-seconds", type=int)
    begin.add_argument("--pretty", action="store_true")

    for name, role in (("append-user", "user"), ("append-assistant", "assistant")):
        command = sub.add_parser(name, help=f"record a {role} message")
        command.add_argument("run_id")
        command.add_argument("--session-id", help="required when the run is not currently readable")
        command.add_argument("--branch-id")
        add_text_input(command)
        command.add_argument("--status", choices=["pending", "streaming", "complete", "partial", "failed"], default="complete")
        command.add_argument("--provider-message-id")
        command.add_argument("--message-id")
        command.add_argument("--idempotency-key")
        command.add_argument("--pretty", action="store_true")

    finish = sub.add_parser("finish-run", help="finish an Agent run")
    finish.add_argument("run_id")
    finish.add_argument("--session-id", help="required when the database is temporarily unavailable")
    finish.add_argument("--status", choices=["completed", "partial", "failed", "cancelled"], required=True)
    finish.add_argument("--error-json")
    finish.add_argument("--pretty", action="store_true")

    checkpoint_command = sub.add_parser("checkpoint", help="create a portable checkpoint")
    checkpoint_command.add_argument("session_id")
    checkpoint_command.add_argument("--summary-json")
    checkpoint_command.add_argument("--summary-file")
    checkpoint_command.add_argument("--created-by")
    checkpoint_command.add_argument("--branch-id")
    checkpoint_command.add_argument("--source-budget-tokens", type=int, default=1800)
    checkpoint_command.add_argument("--pretty", action="store_true")

    branch_create = sub.add_parser("branch-create", help="fork a branch with a copied portable history")
    branch_create.add_argument("session_id")
    branch_create.add_argument("--name")
    branch_create.add_argument("--branch-id")
    branch_create.add_argument("--from-branch-id")
    branch_create.add_argument("--metadata-json")
    branch_create.add_argument("--no-switch", action="store_true")
    branch_create.add_argument("--pretty", action="store_true")

    branch_switch = sub.add_parser("branch-switch", help="switch the active branch")
    branch_switch.add_argument("session_id")
    branch_switch.add_argument("branch_id")
    branch_switch.add_argument("--pretty", action="store_true")

    branch_list = sub.add_parser("branch-list", help="list session branches")
    branch_list.add_argument("session_id")
    branch_list.add_argument("--pretty", action="store_true")

    lease_acquire = sub.add_parser("lease-acquire", help="acquire a branch lease")
    lease_acquire.add_argument("session_id")
    lease_acquire.add_argument("branch_id")
    lease_acquire.add_argument("--agent", required=True)
    lease_acquire.add_argument("--ttl-seconds", type=int, default=1800)
    lease_acquire.add_argument("--pretty", action="store_true")

    lease_release = sub.add_parser("lease-release", help="release a branch lease")
    lease_release.add_argument("session_id")
    lease_release.add_argument("branch_id")
    lease_release.add_argument("--agent")
    lease_release.add_argument("--pretty", action="store_true")

    lease_check = sub.add_parser("lease-status", help="read a branch lease")
    lease_check.add_argument("session_id")
    lease_check.add_argument("branch_id")
    lease_check.add_argument("--pretty", action="store_true")

    for name in ("probe", "resume"):
        command = sub.add_parser(name, help=f"{name} a session with a bounded context budget")
        command.add_argument("session_id")
        command.add_argument("--budget-tokens", type=int, default=1200 if name == "resume" else 600)
        command.add_argument("--pretty", action="store_true")

    expand_command = sub.add_parser("expand", help="append relevant history using a cursor")
    expand_command.add_argument("cursor")
    expand_command.add_argument("--query")
    expand_command.add_argument("--budget-tokens", type=int, default=1600)
    expand_command.add_argument("--allow-inactive-branch", action="store_true")
    expand_command.add_argument("--pretty", action="store_true")

    native = sub.add_parser("save-native-state", help="save optional provider-specific state")
    native.add_argument("session_id")
    native.add_argument("--provider", required=True)
    native.add_argument("--kind", required=True)
    native.add_argument("--payload-json", required=True)
    native.add_argument("--portable", action="store_true")
    native.add_argument("--pretty", action="store_true")

    artifact = sub.add_parser("artifact-add", help="record a file or URI artifact reference")
    artifact.add_argument("session_id")
    artifact.add_argument("--path", required=True)
    artifact.add_argument("--kind", required=True)
    artifact.add_argument("--run-id")
    artifact.add_argument("--branch-id")
    artifact.add_argument("--sha256")
    artifact.add_argument("--size-bytes", type=int)
    artifact.add_argument("--metadata-json")
    artifact.add_argument("--pretty", action="store_true")

    artifact_list = sub.add_parser("artifact-list", help="list artifact references")
    artifact_list.add_argument("session_id")
    artifact_list.add_argument("--branch-id")
    artifact_list.add_argument("--pretty", action="store_true")

    export = sub.add_parser("export", help="export a portable session envelope")
    export.add_argument("session_id")
    export.add_argument("--output")
    export.add_argument("--pretty", action="store_true")

    recover = sub.add_parser("recover", help="replay JSONL spool records into SQLite")
    recover.add_argument("--pretty", action="store_true")

    check = sub.add_parser("verify", help="run database integrity and count checks")
    check.add_argument("--pretty", action="store_true")

    config = sub.add_parser("config", help="show or update local configuration")
    config_actions = config.add_subparsers(dest="action", required=True)
    config_actions.add_parser("show", help="show effective configuration")
    config_set = config_actions.add_parser("set", help="set a dotted configuration key")
    config_set.add_argument("key")
    config_set.add_argument("--value", required=True)
    config_set.add_argument("--pretty", action="store_true")

    backup = sub.add_parser("backup", help="create a consistent SQLite backup")
    backup.add_argument("--output")
    backup.add_argument("--keep", type=int)
    backup.add_argument("--force", action="store_true")
    backup.add_argument("--pretty", action="store_true")

    prune_command = sub.add_parser("prune", help="preview or apply local retention policy")
    prune_command.add_argument("--older-than-days", type=int)
    prune_command.add_argument("--processed-spool-days", type=int)
    prune_command.add_argument("--apply", action="store_true", help="perform deletion; default is dry-run")
    prune_command.add_argument("--pretty", action="store_true")
    return parser


def command_main(args: argparse.Namespace, store: Store) -> Any:
    if args.command == "init":
        store.connect().close()
        if not store.config_path.exists():
            store.write_config(store.read_config())
        return {
            "database": str(store.db_path),
            "home": str(store.home),
            "config": str(store.config_path),
        }
    if args.command == "config":
        return config_command(store, args)
    if args.command == "create":
        return create_session(store, args)
    if args.command == "list":
        with store.connect() as conn:
            rows = conn.execute("SELECT * FROM sessions ORDER BY updated_at DESC").fetchall()
            return [
                {
                    "session_id": row["session_id"],
                    "title": row["title"],
                    "status": row["status"],
                    "active_branch_id": row["active_branch_id"],
                    "updated_at": row["updated_at"],
                }
                for row in rows
            ]
    if args.command == "begin-run":
        return begin_run(store, args)
    if args.command == "append-user":
        return append_message(store, args, "user")
    if args.command == "append-assistant":
        return append_message(store, args, "assistant")
    if args.command == "finish-run":
        return finish_run(store, args)
    if args.command == "checkpoint":
        return checkpoint(store, args)
    if args.command == "branch-create":
        return create_branch(store, args)
    if args.command == "branch-switch":
        return switch_branch(store, args)
    if args.command == "branch-list":
        return list_branches(store, args)
    if args.command == "lease-acquire":
        return acquire_lease(store, args)
    if args.command == "lease-release":
        return release_lease(store, args)
    if args.command == "lease-status":
        return lease_status(store, args)
    if args.command == "probe":
        with store.connect() as conn:
            result = manifest(conn, args.session_id)
        result["mode"] = "probe"
        result["selected_tokens"] = estimate_tokens(result)
        return result
    if args.command == "resume":
        if args.budget_tokens <= 0:
            raise ValueError("budget-tokens must be positive")
        return resume(store, args)
    if args.command == "expand":
        if args.budget_tokens <= 0:
            raise ValueError("budget-tokens must be positive")
        return expand(store, args)
    if args.command == "save-native-state":
        return save_native_state(store, args)
    if args.command == "artifact-add":
        return add_artifact(store, args)
    if args.command == "artifact-list":
        return list_artifacts(store, args)
    if args.command == "export":
        return export_session(store, args)
    if args.command == "recover":
        return recover_spool(store)
    if args.command == "verify":
        return verify(store)
    if args.command == "backup":
        return backup_database(store, args)
    if args.command == "prune":
        return prune(store, args)
    raise ValueError(f"unknown command: {args.command}")


def main(argv: Iterable[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(list(argv) if argv is not None else None)
    store = Store(home=args.home, db_path=args.db)
    try:
        result = command_main(args, store)
        pretty = bool(getattr(args, "pretty", False))
        if not (args.command == "export" and not args.output):
            emit(result, pretty=pretty)
        return 0
    except (OSError, ValueError, sqlite3.Error, json.JSONDecodeError) as exc:
        print(f"{APP_NAME}: error: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())