"""Reader for Cursor's state.vscdb (cursorDiskKV) SQLite databases.

Extracts composer sessions with their bubble entries and agent KV blobs
from the globalStorage/state.vscdb cursorDiskKV table.

Note on NULL values: Cursor uses SQLite's ON CONFLICT REPLACE for its KV tables.
When sessions are deleted or archived, Cursor writes NULL as the value rather than
deleting the row — the key remains as a tombstone. In real DBs ~4% of composerData
rows and ~16% of bubbleId rows have NULL values. All reader functions handle this
by skipping NULLs (there is no data to extract from a deleted session).
"""

from __future__ import annotations

import json
import sqlite3
import zlib
from dataclasses import dataclass, field
from typing import Any, Iterator


@dataclass
class VscdbSession:
    """A single composer session extracted from state.vscdb."""

    composer_id: str
    composer_data: dict[str, Any]
    bubble_entries: dict[str, dict[str, Any]]  # bubbleId:<cid>:<bid> → parsed JSON
    agent_kv_entries: dict[str, bytes]  # agentKv:blob:<hash> → raw bytes
    db_path: str


def iter_sessions(db_path: str) -> Iterator[VscdbSession]:
    """Yield VscdbSession objects for each composerData entry in the database.

    Reads all composerData:<id> rows, then for each composer ID loads
    the corresponding bubbleId:<cid>:* rows. Agent KV blobs are loaded
    once and shared across sessions (content-addressed).
    """
    # Load all composer data rows
    composer_rows = _read_kv_rows(db_path, "composerData:")
    if not composer_rows:
        return

    # Load all agent KV blobs (shared, content-addressed)
    agent_kv = _read_kv_rows(db_path, "agentKv:blob:")

    for key, raw in composer_rows.items():
        # key is "composerData:<composerId>"
        composer_id = key.removeprefix("composerData:")
        if not composer_id:
            continue

        # NULL values are tombstones for deleted/archived sessions — skip them
        parsed = _parse_json(raw)
        if parsed is None or not isinstance(parsed, dict):
            continue

        # Load bubble entries for this composer
        bubble_prefix = f"bubbleId:{composer_id}:"
        bubble_rows = _read_kv_rows(db_path, bubble_prefix)
        bubble_entries: dict[str, dict[str, Any]] = {}
        for bkey, braw in bubble_rows.items():
            bparsed = _parse_json(braw)
            if bparsed is not None and isinstance(bparsed, dict):
                bubble_entries[bkey] = bparsed

        yield VscdbSession(
            composer_id=composer_id,
            composer_data=parsed,
            bubble_entries=bubble_entries,
            agent_kv_entries=agent_kv,
            db_path=db_path,
        )


def scan_session_timestamps(db_path: str) -> dict[str, int]:
    """Quick-scan: return {composerId: lastUpdatedAt} without loading full session data.

    Reads composerData:* keys from cursorDiskKV, parses just enough JSON to extract
    lastUpdatedAt (or createdAt as fallback). Much faster than iter_sessions() since
    it doesn't load bubbles or agent KV blobs.
    """
    result: dict[str, int] = {}
    try:
        conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
        try:
            table = _resolve_table(conn)
            cursor = conn.execute(
                f"SELECT key, value FROM {table} WHERE key LIKE 'composerData:%'",
            )
            for key, val in cursor:
                composer_id = key.removeprefix("composerData:")
                if not composer_id:
                    continue
                # NULL values are tombstones for deleted/archived sessions — skip
                if val is None:
                    continue
                if isinstance(val, str):
                    val = val.encode("utf-8")
                parsed = _parse_json(val)
                if parsed is None or not isinstance(parsed, dict):
                    continue
                ts = parsed.get("lastUpdatedAt") or parsed.get("createdAt") or 0
                if isinstance(ts, (int, float)):
                    result[composer_id] = int(ts)
        finally:
            conn.close()
    except (sqlite3.Error, OSError):
        pass
    return result


def load_session(db_path: str, composer_id: str) -> VscdbSession | None:
    """Load a single session by composer ID.

    More efficient than iter_sessions() for incremental updates since it only
    reads the specific composerData, its bubbleId entries, and relevant agentKv blobs.
    """
    try:
        conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
        try:
            table = _resolve_table(conn)

            # Load composer data
            row = conn.execute(
                f"SELECT value FROM {table} WHERE key = ?",
                (f"composerData:{composer_id}",),
            ).fetchone()
            if row is None:
                return None

            # NULL value = tombstone for deleted session
            val = row[0]
            if isinstance(val, str):
                val = val.encode("utf-8")
            parsed = _parse_json(val)
            if parsed is None or not isinstance(parsed, dict):
                return None

            # Load bubble entries (NULL bubble values are also tombstones — skipped by _parse_json)
            bubble_prefix = f"bubbleId:{composer_id}:"
            bubble_cursor = conn.execute(
                f"SELECT key, value FROM {table} WHERE key LIKE ? || '%'",
                (bubble_prefix,),
            )
            bubble_entries: dict[str, dict[str, Any]] = {}
            for bkey, braw in bubble_cursor:
                if isinstance(braw, str):
                    braw = braw.encode("utf-8")
                bparsed = _parse_json(braw)
                if bparsed is not None and isinstance(bparsed, dict):
                    bubble_entries[bkey] = bparsed

            # Load agent KV blobs referenced by conversationState hashes
            agent_kv_entries: dict[str, bytes] = {}
            conv_state = parsed.get("conversationState")
            if isinstance(conv_state, dict):
                for v in conv_state.values():
                    if isinstance(v, str) and len(v) == 64:  # SHA-256 hash
                        agent_key = f"agentKv:blob:{v}"
                        arow = conn.execute(
                            f"SELECT value FROM {table} WHERE key = ?",
                            (agent_key,),
                        ).fetchone()
                        if arow is not None:
                            aval = arow[0]
                            if isinstance(aval, str):
                                aval = aval.encode("utf-8")
                            agent_kv_entries[agent_key] = aval

            return VscdbSession(
                composer_id=composer_id,
                composer_data=parsed,
                bubble_entries=bubble_entries,
                agent_kv_entries=agent_kv_entries,
                db_path=db_path,
            )
        finally:
            conn.close()
    except (sqlite3.Error, OSError):
        return None


def read_item_table(db_path: str, key: str) -> str | None:
    """Read a single key from the ItemTable and return as string.

    Returns None if the key doesn't exist or on any error.
    """
    try:
        conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
        try:
            row = conn.execute(
                "SELECT value FROM ItemTable WHERE key = ?", (key,)
            ).fetchone()
            if row is None:
                return None
            val = row[0]
            if isinstance(val, bytes):
                return _decompress(val).decode("utf-8", errors="replace")
            return str(val)
        finally:
            conn.close()
    except (sqlite3.Error, OSError):
        return None


def _read_kv_rows(db_path: str, prefix: str) -> dict[str, bytes]:
    """Read all rows from cursorDiskKV whose key starts with prefix.

    Falls back to ItemTable if cursorDiskKV doesn't exist.
    Returns {key: raw_value_bytes} dict.
    """
    result: dict[str, bytes] = {}
    try:
        conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
        try:
            table = _resolve_table(conn)
            cursor = conn.execute(
                f"SELECT key, value FROM {table} WHERE key LIKE ? || '%'",
                (prefix,),
            )
            for key, val in cursor:
                # Skip NULL tombstones (deleted sessions/bubbles)
                if val is None:
                    continue
                if isinstance(val, str):
                    val = val.encode("utf-8")
                result[key] = val
        finally:
            conn.close()
    except (sqlite3.Error, OSError):
        pass
    return result


def _resolve_table(conn: sqlite3.Connection) -> str:
    """Return 'cursorDiskKV' if it exists and has data, else 'ItemTable'."""
    try:
        row = conn.execute(
            "SELECT name FROM sqlite_master WHERE type='table' AND name='cursorDiskKV'"
        ).fetchone()
        if row:
            return "cursorDiskKV"
    except sqlite3.Error:
        pass
    return "ItemTable"


def _decompress(raw: bytes) -> bytes:
    """Decompress bytes if zlib-compressed (0x78 0x9c header), otherwise return as-is."""
    if len(raw) >= 2 and raw[0] == 0x78 and raw[1] in (0x01, 0x5E, 0x9C, 0xDA):
        try:
            return zlib.decompress(raw)
        except zlib.error:
            pass
    return raw


def _parse_json(raw: bytes | None) -> Any | None:
    """Decompress (if needed) and parse JSON. Returns None on failure or NULL input.

    NULL values occur for tombstoned (deleted/archived) rows in cursorDiskKV.
    """
    if raw is None:
        return None
    try:
        data = _decompress(raw)
        text = data.decode("utf-8", errors="replace")
        return json.loads(text)
    except (json.JSONDecodeError, UnicodeDecodeError, ValueError):
        return None
