#!/usr/bin/env python3
"""One-time cleanup script for all Omni-Cortex project databases.

Finds all .omni-cortex/cortex.db files under D:/Projects/ and D:/Clients/,
plus ~/.omni-cortex/global.db, applies retention rules, VACUUMs, and reports
before/after sizes.

Usage:
    python scripts/cleanup_all_dbs.py [--dry-run]
"""

import argparse
import os
import sqlite3
import time
from datetime import datetime, timedelta, timezone
from pathlib import Path

ACTIVITY_CAP = 7000
MESSAGE_CAP = 2000
MAX_AGE_DAYS = 14

SEARCH_ROOTS = [
    Path("D:/Projects"),
    Path("D:/Clients"),
]
GLOBAL_DB = Path.home() / ".omni-cortex" / "global.db"


def find_all_dbs() -> list[Path]:
    """Find all cortex.db files under search roots + global DB."""
    dbs = []

    for root in SEARCH_ROOTS:
        if not root.exists():
            continue
        for dirpath, dirnames, filenames in os.walk(root):
            # Skip nested .omni-cortex dirs (don't recurse into them)
            if ".omni-cortex" in dirnames:
                db_path = Path(dirpath) / ".omni-cortex" / "cortex.db"
                if db_path.exists():
                    dbs.append(db_path)
                dirnames.remove(".omni-cortex")
            # Skip common non-project dirs for speed
            for skip in ["node_modules", ".git", "__pycache__", ".venv", "venv"]:
                if skip in dirnames:
                    dirnames.remove(skip)

    if GLOBAL_DB.exists():
        dbs.append(GLOBAL_DB)

    return sorted(dbs)


def get_file_size_mb(path: Path) -> float:
    """Get file size in MB."""
    return path.stat().st_size / (1024 * 1024)


def get_table_count(conn: sqlite3.Connection, table: str) -> int:
    """Get row count for a table, returning 0 if table doesn't exist."""
    try:
        return conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0]
    except sqlite3.OperationalError:
        return 0


def prune_database(db_path: Path, dry_run: bool = False) -> dict:
    """Apply retention rules to a single database.

    Returns dict with before/after stats.
    """
    size_before = get_file_size_mb(db_path)
    project_name = db_path.parent.parent.name if db_path.name == "cortex.db" else "global"

    conn = sqlite3.connect(str(db_path))
    conn.execute("PRAGMA journal_mode = WAL")
    conn.execute("PRAGMA busy_timeout = 30000")

    activities_before = get_table_count(conn, "activities")
    messages_before = get_table_count(conn, "user_messages")
    memories_count = get_table_count(conn, "memories")

    cutoff = (datetime.now(timezone.utc) - timedelta(days=MAX_AGE_DAYS)).isoformat()

    activities_deleted = 0
    messages_deleted = 0

    if not dry_run:
        retry_backoff = [0.2, 0.5, 1.0, 2.0, 4.0]

        for attempt in range(5):
            try:
                # Prune activities beyond cap
                if activities_before > ACTIVITY_CAP:
                    conn.execute("""
                        DELETE FROM activities WHERE id NOT IN (
                            SELECT id FROM activities ORDER BY timestamp DESC LIMIT ?
                        )
                    """, (ACTIVITY_CAP,))

                # Prune activities older than MAX_AGE_DAYS
                conn.execute("DELETE FROM activities WHERE timestamp < ?", (cutoff,))

                # Prune user_messages
                try:
                    if messages_before > MESSAGE_CAP:
                        conn.execute("""
                            DELETE FROM user_messages WHERE id NOT IN (
                                SELECT id FROM user_messages ORDER BY timestamp DESC LIMIT ?
                            )
                        """, (MESSAGE_CAP,))
                    conn.execute("DELETE FROM user_messages WHERE timestamp < ?", (cutoff,))
                except sqlite3.OperationalError:
                    pass  # Table may not exist

                conn.commit()
                break
            except sqlite3.OperationalError as e:
                if "database is locked" in str(e) and attempt < 4:
                    time.sleep(retry_backoff[attempt])
                    continue
                print(f"  WARNING: Could not prune {project_name}: {e}")
                conn.close()
                return {
                    "project": project_name,
                    "path": str(db_path),
                    "error": str(e),
                }

        activities_after = get_table_count(conn, "activities")
        messages_after = get_table_count(conn, "user_messages")
        activities_deleted = activities_before - activities_after
        messages_deleted = messages_before - messages_after

        # VACUUM to reclaim space
        for attempt in range(5):
            try:
                conn.execute("VACUUM")
                break
            except sqlite3.OperationalError as e:
                if "database is locked" in str(e) and attempt < 4:
                    time.sleep(retry_backoff[attempt])
                    continue
                print(f"  WARNING: VACUUM failed for {project_name}: {e}")
                break

    conn.close()

    size_after = get_file_size_mb(db_path) if not dry_run else size_before

    return {
        "project": project_name,
        "path": str(db_path),
        "size_before_mb": size_before,
        "size_after_mb": size_after,
        "activities_before": activities_before,
        "activities_deleted": activities_deleted,
        "messages_before": messages_before,
        "messages_deleted": messages_deleted,
        "memories": memories_count,
    }


def main():
    parser = argparse.ArgumentParser(description="Clean up all Omni-Cortex databases")
    parser.add_argument("--dry-run", action="store_true", help="Show what would be cleaned without making changes")
    args = parser.parse_args()

    print("=" * 70)
    print("Omni-Cortex Database Cleanup")
    print(f"Retention: {ACTIVITY_CAP} activities, {MESSAGE_CAP} messages, {MAX_AGE_DAYS} days max age")
    if args.dry_run:
        print("MODE: DRY RUN (no changes will be made)")
    print("=" * 70)
    print()

    dbs = find_all_dbs()
    if not dbs:
        print("No databases found.")
        return

    print(f"Found {len(dbs)} database(s):\n")

    results = []
    total_saved = 0.0

    for db_path in dbs:
        project = db_path.parent.parent.name if db_path.name == "cortex.db" else "global"
        print(f"Processing: {project} ({db_path})")

        result = prune_database(db_path, dry_run=args.dry_run)
        results.append(result)

        if "error" in result:
            print(f"  ERROR: {result['error']}\n")
            continue

        saved = result["size_before_mb"] - result["size_after_mb"]
        total_saved += saved

        print(f"  Size: {result['size_before_mb']:.1f} MB -> {result['size_after_mb']:.1f} MB (saved {saved:.1f} MB)")
        print(f"  Activities: {result['activities_before']} -> {result['activities_before'] - result['activities_deleted']} (deleted {result['activities_deleted']})")
        print(f"  Messages: {result['messages_before']} -> {result['messages_before'] - result['messages_deleted']} (deleted {result['messages_deleted']})")
        print(f"  Memories: {result['memories']} (untouched)")
        print()

    # Summary table
    print("=" * 70)
    print("SUMMARY")
    print("=" * 70)
    print(f"{'Project':<25} {'Before':>8} {'After':>8} {'Saved':>8} {'Activities':>12} {'Messages':>10}")
    print("-" * 70)

    total_before = 0.0
    total_after = 0.0

    for r in results:
        if "error" in r:
            print(f"{r['project']:<25} {'ERROR':>8}")
            continue

        before = r["size_before_mb"]
        after = r["size_after_mb"]
        saved = before - after
        total_before += before
        total_after += after

        print(
            f"{r['project']:<25} "
            f"{before:>7.1f}M "
            f"{after:>7.1f}M "
            f"{saved:>7.1f}M "
            f"{r['activities_deleted']:>12} "
            f"{r['messages_deleted']:>10}"
        )

    print("-" * 70)
    print(
        f"{'TOTAL':<25} "
        f"{total_before:>7.1f}M "
        f"{total_after:>7.1f}M "
        f"{total_saved:>7.1f}M"
    )
    print()


if __name__ == "__main__":
    main()
