#!/usr/bin/env python3
"""Manage the video-edit-planner SQLite index.

The index file lives next to the video as <video_stem>.vedit.db.
It tracks transcriptions, clips, and frames to avoid duplicate processing.

Usage:
    python manage_index.py init --video VIDEO_PATH
    python manage_index.py add-transcription --json-path PATH --track-index N --duration S
    python manage_index.py get-transcription
    python manage_index.py add-clip --start S --end E --path PATH --reason "..."
    python manage_index.py add-frames --clip-id N --frames-dir DIR --method scene
    python manage_index.py list
    python manage_index.py check-range --start S --end E
"""

from __future__ import annotations

import argparse
import json
import sqlite3
import sys
import time
from pathlib import Path

SCHEMA = """
CREATE TABLE IF NOT EXISTS transcription (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    json_path TEXT NOT NULL,
    track_index INTEGER NOT NULL,
    created_at REAL NOT NULL,
    duration_s REAL
);

CREATE TABLE IF NOT EXISTS clips (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    start_time REAL NOT NULL,
    end_time REAL NOT NULL,
    path TEXT NOT NULL,
    created_at REAL NOT NULL,
    reason TEXT
);

CREATE TABLE IF NOT EXISTS frames (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    clip_id INTEGER NOT NULL REFERENCES clips(id),
    timestamp REAL NOT NULL,
    path TEXT NOT NULL,
    scene_score REAL,
    method TEXT NOT NULL,
    created_at REAL NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_frames_clip_id ON frames(clip_id);
CREATE INDEX IF NOT EXISTS idx_clips_time ON clips(start_time, end_time);
"""


def db_path_for_video(video_path: Path) -> Path:
    """Return the index database path for a given video file."""
    return video_path.parent / f"{video_path.stem}.vedit.db"


def connect(video_path: Path) -> sqlite3.Connection:
    db = db_path_for_video(video_path)
    if not db.exists():
        sys.exit(f"Index not found: {db}. Run 'init' first.")
    conn = sqlite3.connect(str(db))
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA foreign_keys = ON")
    return conn


def cmd_init(args: argparse.Namespace) -> None:
    video = Path(args.video).expanduser().resolve()
    db = db_path_for_video(video)
    conn = sqlite3.connect(str(db))
    conn.executescript(SCHEMA)
    conn.close()
    print(json.dumps({"status": "ok", "db_path": str(db)}))


def cmd_add_transcription(args: argparse.Namespace) -> None:
    video = Path(args.video).expanduser().resolve()
    conn = connect(video)
    conn.execute(
        "INSERT INTO transcription (json_path, track_index, created_at, duration_s) VALUES (?, ?, ?, ?)",
        (args.json_path, args.track_index, time.time(), args.duration),
    )
    conn.commit()
    cur = conn.execute("SELECT last_insert_rowid()")
    row_id = cur.fetchone()[0]
    conn.close()
    print(json.dumps({"status": "ok", "id": row_id}))


def cmd_get_transcription(args: argparse.Namespace) -> None:
    video = Path(args.video).expanduser().resolve()
    conn = connect(video)
    if args.track is not None:
        rows = conn.execute(
            "SELECT * FROM transcription WHERE track_index = ? ORDER BY created_at DESC",
            (args.track,),
        ).fetchall()
    else:
        rows = conn.execute("SELECT * FROM transcription ORDER BY created_at DESC").fetchall()
    conn.close()
    if not rows:
        print(json.dumps({"found": False}))
        return
    result = [dict(r) for r in rows]
    print(json.dumps({"found": True, "records": result}, ensure_ascii=False))


def cmd_add_clip(args: argparse.Namespace) -> None:
    video = Path(args.video).expanduser().resolve()
    conn = connect(video)
    cur = conn.execute(
        "INSERT INTO clips (start_time, end_time, path, created_at, reason) VALUES (?, ?, ?, ?, ?)",
        (args.start, args.end, args.path, time.time(), args.reason),
    )
    conn.commit()
    clip_id = cur.lastrowid
    conn.close()
    print(json.dumps({"status": "ok", "clip_id": clip_id}))


def cmd_add_frames(args: argparse.Namespace) -> None:
    """Add frame records for all images in a directory, associating with a clip."""
    video = Path(args.video).expanduser().resolve()
    frames_dir = Path(args.frames_dir).expanduser().resolve()
    if not frames_dir.is_dir():
        sys.exit(f"Frames directory not found: {frames_dir}")

    conn = connect(video)

    # Read scene scores from showinfo log if available
    scene_scores: dict[str, float] = {}
    log_file = frames_dir / "showinfo.log"
    if log_file.exists():
        for line in log_file.read_text(errors="replace").splitlines():
            # Parse lavfi.scd.score and lavfi.scd.time from showinfo output
            score = None
            t = None
            for part in line.split():
                if part.startswith("score:"):
                    score = float(part.split(":")[1])
                elif part.startswith("pts_time:"):
                    t = float(part.split(":")[1])
            if score is not None and t is not None:
                scene_scores[f"{t:.3f}"] = score

    # Parse timestamps from filenames: frame_XXXX.jpg
    # The agent should pass --fps and --start to compute timestamps
    fps = args.fps if args.fps else 0.5
    start_offset = args.start if args.start else 0.0

    frame_files = sorted(frames_dir.glob("frame_*.jpg"))
    if not frame_files:
        # Try other extensions
        frame_files = sorted(frames_dir.glob("frame_*.png"))

    if not frame_files:
        sys.exit(f"No frame files found in {frames_dir}")

    count = 0
    for i, frame_path in enumerate(frame_files):
        # Estimate timestamp: frame number * (1/fps) + start_offset
        # This is approximate; the agent can refine with --timestamps
        timestamp = start_offset + (i / fps)

        # Try to match with scene score
        scene_score = scene_scores.get(f"{timestamp:.3f}")

        conn.execute(
            "INSERT INTO frames (clip_id, timestamp, path, scene_score, method, created_at) VALUES (?, ?, ?, ?, ?, ?)",
            (args.clip_id, timestamp, str(frame_path), scene_score, args.method, time.time()),
        )
        count += 1

    conn.commit()
    conn.close()
    print(json.dumps({"status": "ok", "frames_added": count}))


def cmd_list(args: argparse.Namespace) -> None:
    video = Path(args.video).expanduser().resolve()
    conn = connect(video)

    transcriptions = conn.execute("SELECT * FROM transcription ORDER BY created_at DESC").fetchall()
    clips = conn.execute("SELECT * FROM clips ORDER BY start_time").fetchall()

    result = {"transcriptions": [dict(r) for r in transcriptions], "clips": []}
    for clip in clips:
        clip_dict = dict(clip)
        frames = conn.execute(
            "SELECT * FROM frames WHERE clip_id = ? ORDER BY timestamp", (clip["id"],)
        ).fetchall()
        clip_dict["frames"] = [dict(f) for f in frames]
        result["clips"].append(clip_dict)

    conn.close()
    print(json.dumps(result, ensure_ascii=False, indent=2))


def cmd_check_range(args: argparse.Namespace) -> None:
    """Check if a time range overlaps with any existing clips."""
    video = Path(args.video).expanduser().resolve()
    conn = connect(video)
    # Find clips that overlap with the requested range
    rows = conn.execute(
        """SELECT * FROM clips
           WHERE start_time < ? AND end_time > ?
           ORDER BY start_time""",
        (args.end, args.start),
    ).fetchall()
    conn.close()
    if not rows:
        print(json.dumps({"found": False}))
    else:
        print(json.dumps({"found": True, "clips": [dict(r) for r in rows]}, ensure_ascii=False))


def cmd_clean(args: argparse.Namespace) -> None:
    """Remove frame and clip records for files that no longer exist on disk."""
    video = Path(args.video).expanduser().resolve()
    conn = connect(video)

    # Check frames
    frame_rows = conn.execute("SELECT id, path FROM frames").fetchall()
    frames_removed = 0
    for row in frame_rows:
        if not Path(row["path"]).exists():
            conn.execute("DELETE FROM frames WHERE id = ?", (row["id"],))
            frames_removed += 1

    # Check clips
    clip_rows = conn.execute("SELECT id, path FROM clips").fetchall()
    clips_removed = 0
    for row in clip_rows:
        if not Path(row["path"]).exists():
            conn.execute("DELETE FROM frames WHERE clip_id = ?", (row["id"],))
            conn.execute("DELETE FROM clips WHERE id = ?", (row["id"],))
            clips_removed += 1

    conn.commit()
    conn.close()
    print(json.dumps({"status": "ok", "frames_removed": frames_removed, "clips_removed": clips_removed}))


def main() -> None:
    parser = argparse.ArgumentParser(description="Manage video-edit-planner index database.")
    sub = parser.add_subparsers(dest="command", required=True)

    # init
    p_init = sub.add_parser("init", help="Initialize index database")
    p_init.add_argument("--video", required=True, help="Path to the video file")

    # add-transcription
    p_at = sub.add_parser("add-transcription", help="Add a transcription record")
    p_at.add_argument("--video", required=True)
    p_at.add_argument("--json-path", required=True)
    p_at.add_argument("--track-index", required=True, type=int)
    p_at.add_argument("--duration", type=float, default=None)

    # get-transcription
    p_gt = sub.add_parser("get-transcription", help="Get transcription records (all or by track)")
    p_gt.add_argument("--video", required=True)
    p_gt.add_argument("--track", type=int, default=None, help="Filter by track index (omit to get all)")

    # add-clip
    p_ac = sub.add_parser("add-clip", help="Add a clip record")
    p_ac.add_argument("--video", required=True)
    p_ac.add_argument("--start", required=True, type=float, help="Start time in seconds")
    p_ac.add_argument("--end", required=True, type=float, help="End time in seconds")
    p_ac.add_argument("--path", required=True, help="Path to the clip file")
    p_ac.add_argument("--reason", default=None)

    # add-frames
    p_af = sub.add_parser("add-frames", help="Add frame records from a directory")
    p_af.add_argument("--video", required=True)
    p_af.add_argument("--clip-id", required=True, type=int)
    p_af.add_argument("--frames-dir", required=True)
    p_af.add_argument("--method", required=True, choices=["scene", "sample", "both"])
    p_af.add_argument("--fps", type=float, default=0.5, help="Sampling fps for timestamp estimation")
    p_af.add_argument("--start", type=float, default=0.0, help="Start offset in seconds")

    # list
    p_list = sub.add_parser("list", help="List all records")
    p_list.add_argument("--video", required=True)

    # check-range
    p_cr = sub.add_parser("check-range", help="Check if a time range has existing clips")
    p_cr.add_argument("--video", required=True)
    p_cr.add_argument("--start", required=True, type=float)
    p_cr.add_argument("--end", required=True, type=float)

    # clean
    p_clean = sub.add_parser("clean", help="Remove records for files that no longer exist")
    p_clean.add_argument("--video", required=True)

    args = parser.parse_args()

    commands = {
        "init": cmd_init,
        "add-transcription": cmd_add_transcription,
        "get-transcription": cmd_get_transcription,
        "add-clip": cmd_add_clip,
        "add-frames": cmd_add_frames,
        "list": cmd_list,
        "check-range": cmd_check_range,
        "clean": cmd_clean,
    }
    commands[args.command](args)


if __name__ == "__main__":
    main()