#!/usr/bin/env python3
"""Export Cursor chat history for ~/Documents/repos workspaces into this archive."""

from __future__ import annotations

import hashlib
import json
import re
import shutil
import sqlite3
import subprocess
from datetime import datetime, timezone
from pathlib import Path
from urllib.parse import unquote, urlparse

ARCHIVE_ROOT = Path(__file__).resolve().parents[1]
OLD_HOME = Path.home()
OLD_REPOS_ROOT = OLD_HOME / "Documents" / "repos"
CURSOR_PROJECTS = OLD_HOME / ".cursor" / "projects"
CURSOR_USER = OLD_HOME / ".config" / "Cursor" / "User"
GLOBAL_DB = CURSOR_USER / "globalStorage" / "state.vscdb"
STORAGE_JSON = CURSOR_USER / "globalStorage" / "storage.json"
IDE_STATE = OLD_HOME / ".cursor" / "ide_state.json"
WS_ROOT = CURSOR_USER / "workspaceStorage"


def folder_uri(path: Path) -> str:
    return f"file://{path.resolve()}"


def workspace_id(path: Path) -> str:
    return hashlib.md5(folder_uri(path).encode()).hexdigest()


def path_to_slug(path: Path) -> str:
    return str(path.resolve()).strip("/").replace("/", "-")


def repo_name_from_path(path: str) -> str | None:
    marker = "/Documents/repos/"
    if marker not in path:
        return None
    return path.split(marker, 1)[1].split("/")[0]


def sqlite_backup(src: Path, dst: Path) -> None:
    dst.parent.mkdir(parents=True, exist_ok=True)
    if dst.exists():
        dst.unlink()
    src_conn = sqlite3.connect(f"file:{src}?mode=ro", uri=True)
    dst_conn = sqlite3.connect(dst)
    src_conn.backup(dst_conn)
    src_conn.close()
    dst_conn.close()


def parse_folder_uri(uri: str) -> str:
    if uri.startswith("file://"):
        parsed = urlparse(uri)
        return unquote(parsed.path)
    return uri


def collect_composer_index(conn: sqlite3.Connection) -> dict[str, list[str]]:
    """Map repo_name -> list of composerIds."""
    by_repo: dict[str, list[str]] = {}
    row = conn.execute(
        "SELECT value FROM ItemTable WHERE key='composer.composerHeaders'"
    ).fetchone()
    if not row:
        return by_repo
    data = json.loads(row[0])
    for composer in data.get("allComposers", []):
        cid = composer.get("composerId")
        if not cid:
            continue
        wi = composer.get("workspaceIdentifier") or {}
        uri = wi.get("uri") or {}
        fs_path = uri.get("fsPath") or uri.get("path") or ""
        name = repo_name_from_path(fs_path)
        if name:
            by_repo.setdefault(name, []).append(cid)
    return by_repo


def export_composer_blobs(conn: sqlite3.Connection, composer_ids: list[str], dest: Path) -> int:
    dest.mkdir(parents=True, exist_ok=True)
    count = 0
    for cid in composer_ids:
        key = f"composerData:{cid}"
        row = conn.execute(
            "SELECT value FROM cursorDiskKV WHERE key=?", (key,)
        ).fetchone()
        if not row:
            continue
        out = dest / f"{cid}.json"
        try:
            parsed = json.loads(row[0])
            out.write_text(
                json.dumps(parsed, ensure_ascii=False, indent=2) + "\n",
                encoding="utf-8",
            )
        except json.JSONDecodeError:
            out.write_text(row[0], encoding="utf-8")
        count += 1
    return count


def copy_tree(src: Path, dst: Path) -> bool:
    if not src.exists():
        return False
    if dst.exists():
        shutil.rmtree(dst)
    shutil.copytree(src, dst, dirs_exist_ok=True)
    return True


def discover_repos_from_workspace_storage() -> dict[str, dict]:
    repos: dict[str, dict] = {}
    for ws_dir in sorted(WS_ROOT.iterdir()):
        wj = ws_dir / "workspace.json"
        if not wj.is_file():
            continue
        data = json.loads(wj.read_text(encoding="utf-8"))
        folder = data.get("folder", "")
        if "/Documents/repos/" not in folder:
            continue
        abs_path = parse_folder_uri(folder)
        name = repo_name_from_path(abs_path)
        if not name:
            continue
        entry = repos.setdefault(
            name,
            {
                "repo_name": name,
                "old_absolute_path": abs_path,
                "old_folder_uri": folder,
                "old_workspace_ids": [],
                "old_project_slug": path_to_slug(Path(abs_path)),
                "workspace_storage_sizes": {},
            },
        )
        ws_id = ws_dir.name
        if ws_id not in entry["old_workspace_ids"]:
            entry["old_workspace_ids"].append(ws_id)
        size = sum(f.stat().st_size for f in ws_dir.rglob("*") if f.is_file())
        entry["workspace_storage_sizes"][ws_id] = size
    return repos


def discover_project_slugs() -> dict[str, str]:
    """slug -> repo_name (if mappable)."""
    mapping: dict[str, str] = {}
    prefix = "home-ruslanpi-Documents-repos-"
    for proj in CURSOR_PROJECTS.iterdir():
        if not proj.is_dir():
            continue
        name = proj.name
        if name == "home-ruslanpi-Documents-repos":
            mapping[name] = "__repos_root__"
        elif name.startswith(prefix):
            mapping[name] = name[len(prefix) :]
        elif "Documents-repos-lmru-devsecops-ansible" in name:
            mapping[name] = "lmru--devsecops--ansible"
    return mapping


def pick_primary_workspace_id(entry: dict) -> str:
    sizes = entry.get("workspace_storage_sizes") or {}
    if not sizes:
        return entry["old_workspace_ids"][0]
    return max(sizes, key=sizes.get)


def main() -> None:
    print("Archive root:", ARCHIVE_ROOT)
    print("Exporting from:", OLD_REPOS_ROOT)

    repos = discover_repos_from_workspace_storage()
    slug_map = discover_project_slugs()

    for slug, repo_name in slug_map.items():
        if repo_name == "__repos_root__":
            continue
        if repo_name not in repos:
            repos[repo_name] = {
                "repo_name": repo_name,
                "old_absolute_path": str(OLD_REPOS_ROOT / repo_name),
                "old_folder_uri": folder_uri(OLD_REPOS_ROOT / repo_name),
                "old_workspace_ids": [],
                "old_project_slug": slug,
                "workspace_storage_sizes": {},
            }

    # Global backup
    global_dir = ARCHIVE_ROOT / "global"
    global_dir.mkdir(parents=True, exist_ok=True)
    if GLOBAL_DB.exists():
        print("Backing up global state.vscdb ...")
        sqlite_backup(GLOBAL_DB, global_dir / "state.vscdb")
    if STORAGE_JSON.exists():
        shutil.copy2(STORAGE_JSON, global_dir / "storage.json")
    if IDE_STATE.exists():
        shutil.copy2(IDE_STATE, global_dir / "ide_state.json")

    composer_by_repo: dict[str, list[str]] = {}
    if GLOBAL_DB.exists():
        conn = sqlite3.connect(f"file:{GLOBAL_DB}?mode=ro", uri=True)
        composer_by_repo = collect_composer_index(conn)

    stats = {
        "repos_count": 0,
        "agent_transcript_files": 0,
        "workspace_storage_dirs": 0,
        "composer_json_exports": 0,
    }

    for repo_name, entry in sorted(repos.items()):
        stats["repos_count"] += 1
        repo_dir = ARCHIVE_ROOT / "repos" / repo_name
        repo_dir.mkdir(parents=True, exist_ok=True)

        old_path = Path(entry["old_absolute_path"])
        entry["expected_workspace_id"] = workspace_id(old_path)
        entry["old_project_slug"] = path_to_slug(old_path)
        entry["composer_ids"] = sorted(set(composer_by_repo.get(repo_name, [])))
        entry["primary_workspace_id"] = (
            pick_primary_workspace_id(entry)
            if entry.get("old_workspace_ids")
            else entry["expected_workspace_id"]
        )

        # agent-transcripts
        slug = entry["old_project_slug"]
        proj_dir = CURSOR_PROJECTS / slug
        at_src = proj_dir / "agent-transcripts"
        if copy_tree(at_src, repo_dir / "agent-transcripts"):
            stats["agent_transcript_files"] += len(
                list((repo_dir / "agent-transcripts").rglob("*.jsonl"))
            )
            entry["has_agent_transcripts"] = True
        else:
            entry["has_agent_transcripts"] = False

        # extra slugs (renamed repos)
        extra_slugs = [
            s for s, r in slug_map.items() if r == repo_name and s != slug
        ]
        entry["extra_project_slugs"] = extra_slugs
        for extra in extra_slugs:
            extra_at = CURSOR_PROJECTS / extra / "agent-transcripts"
            extra_dst = repo_dir / "agent-transcripts-extra" / extra
            if copy_tree(extra_at, extra_dst):
                stats["agent_transcript_files"] += len(
                    list(extra_dst.rglob("*.jsonl"))
                )

        # workspace-storage (all ids, primary highlighted)
        ws_archive = repo_dir / "workspace-storage"
        ws_archive.mkdir(exist_ok=True)
        entry["has_workspace_storage"] = False
        for ws_id in entry.get("old_workspace_ids") or []:
            src = WS_ROOT / ws_id
            if copy_tree(src, ws_archive / ws_id):
                stats["workspace_storage_dirs"] += 1
                entry["has_workspace_storage"] = True

        # composer exports
        if entry["composer_ids"] and GLOBAL_DB.exists():
            conn = sqlite3.connect(f"file:{GLOBAL_DB}?mode=ro", uri=True)
            n = export_composer_blobs(
                conn, entry["composer_ids"], repo_dir / "composers"
            )
            stats["composer_json_exports"] += n
            conn.close()

        (repo_dir / "meta.json").write_text(
            json.dumps(entry, ensure_ascii=False, indent=2) + "\n",
            encoding="utf-8",
        )

    # orphan slugs (transcripts without repo folder on disk)
    orphans = []
    for slug, repo_name in slug_map.items():
        if repo_name == "__repos_root__":
            orphans.append({"slug": slug, "kind": "repos_root"})
            proj = CURSOR_PROJECTS / slug
            if proj.exists():
                dst = ARCHIVE_ROOT / "orphans" / slug
                copy_tree(proj, dst)
            continue
        if repo_name not in repos:
            orphans.append({"slug": slug, "repo_name": repo_name})
            dst = ARCHIVE_ROOT / "orphans" / slug
            copy_tree(CURSOR_PROJECTS / slug, dst)

    source_profile = {
        "exported_at": datetime.now(timezone.utc).isoformat(),
        "old_home": str(OLD_HOME),
        "old_repos_root": str(OLD_REPOS_ROOT),
        "old_username": OLD_HOME.name,
        "workspace_id_uri_format": "file://{absolute_path}",
        "cursor_projects_dir": str(CURSOR_PROJECTS),
        "cursor_user_dir": str(CURSOR_USER),
    }

    manifest = {
        "version": 1,
        "source": source_profile,
        "stats": stats,
        "repos": repos,
        "orphan_project_slugs": orphans,
    }

    (ARCHIVE_ROOT / "source-profile.json").write_text(
        json.dumps(source_profile, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    (ARCHIVE_ROOT / "manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )

    print(json.dumps(stats, indent=2))
    print("Done. Files written to", ARCHIVE_ROOT)


if __name__ == "__main__":
    main()