#!/usr/bin/env python3
"""
VLM-based semantic search for video keyframes.
Architecture:
  1. Extract keyframes from video at configurable interval
  2. Send each keyframe to VLM (Ollama LLaVA) for description
  3. Store frame descriptions in SQLite with embedding
  4. Search via text similarity using embedding model

Usage:
  python3 vlm_search.py index <file_uuid>    # Index a video
  python3 vlm_search.py search <query>        # Search indexed frames
  python3 vlm_search.py status                # Show index status
  python3 vlm_search.py benchmark <file_uuid> # Benchmark VLM speed
"""
import sys, os, json, time, sqlite3, argparse, base64, io, struct
from pathlib import Path

import urllib.request
import urllib.parse

# ─── Configuration ───
CORE_API = "http://localhost:3002"
PROXY = "http://localhost:8888"
API_KEY = "muser_68600856036340bcafc01930eb4bd839_1774418104_97221b69"
DB_PATH = Path(__file__).parent / "data" / "vlm_index.sqlite"
FRAME_INTERVAL = 120  # every N frames (default ~5s at 24fps)
VLM_MODEL = "llava"
EMBED_MODEL = "mxbai-embed-large"

def api_get(path):
    url = f"{CORE_API}{path}?api_key={API_KEY}"
    with urllib.request.urlopen(url) as r:
        return json.loads(r.read())

def api_post(path, body):
    url = f"{CORE_API}{path}?api_key={API_KEY}"
    data = json.dumps(body).encode()
    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req) as r:
        return json.loads(r.read())

def get_thumbnail_proxy(file_uuid, frame):
    url = f"{PROXY}/api/v1/file/{file_uuid}/thumbnail?frame={frame}"
    with urllib.request.urlopen(url) as r:
        return r.read()

# ─── SQLite Index ───
def init_db():
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(str(DB_PATH))
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("""
        CREATE TABLE IF NOT EXISTS frames (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            file_uuid TEXT NOT NULL,
            frame INTEGER NOT NULL,
            time_sec REAL,
            description TEXT,
            embedding BLOB,
            created_at REAL DEFAULT (julianday('now'))
        )""")
    conn.execute("CREATE INDEX IF NOT EXISTS idx_frames_desc ON frames(description)")
    conn.execute("CREATE INDEX IF NOT EXISTS idx_frames_file ON frames(file_uuid)")
    conn.commit()
    return conn

def get_embedding(text):
    """Use Ollama embedding model to get vector."""
    payload = json.dumps({"model": EMBED_MODEL, "prompt": text}).encode()
    req = urllib.request.Request(
        "http://localhost:11434/api/embeddings",
        data=payload,
        headers={"Content-Type": "application/json"}
    )
    try:
        with urllib.request.urlopen(req, timeout=30) as r:
            data = json.loads(r.read())
            return data.get("embedding")
    except Exception as e:
        print(f"  Embedding error: {e}")
        return None

def cosine_similarity(a, b):
    dot = sum(x*y for x,y in zip(a,b))
    na = sum(x*x for x in a)**0.5
    nb = sum(x*x for x in b)**0.5
    return dot / (na * nb) if na and nb else 0

# ─── VLM Inference ───
def vlm_describe(image_bytes, prompt="Describe this image in 1-2 sentences. Focus on what is happening, people, objects, and scene."):
    b64 = base64.b64encode(image_bytes).decode()
    payload = json.dumps({
        "model": VLM_MODEL,
        "prompt": prompt,
        "images": [b64],
        "stream": False,
        "options": {"num_predict": 128}
    }).encode()
    req = urllib.request.Request(
        "http://localhost:11434/api/generate",
        data=payload,
        headers={"Content-Type": "application/json"}
    )
    with urllib.request.urlopen(req, timeout=120) as r:
        data = json.loads(r.read())
    return data.get("response", "").strip()

# ─── Commands ───
def cmd_index(file_uuid, interval=FRAME_INTERVAL, limit=None):
    conn = init_db()
    info = api_get(f"/api/v1/file/{file_uuid}")
    if not info.get("file_uuid"):
        print(f"File not found: {file_uuid}")
        return
    file_name = info.get("file_name", "?")
    fps = info.get("fps", 24)
    duration = info.get("duration", 0)
    total = 0
    frame = 0
    times = []
    print(f"Indexing: {file_name} ({duration:.1f}s @ {fps:.2f}fps)")
    while True:
        if limit and total >= limit:
            break
        try:
            img = get_thumbnail_proxy(file_uuid, frame)
        except Exception as e:
            if "404" in str(e):
                break  # no more frames
            print(f"  Error frame {frame}: {e}")
            frame += interval
            continue
        t0 = time.time()
        try:
            desc = vlm_describe(img, "Describe this image briefly. What are people doing? What are they wearing? What objects/scene?")
        except Exception as e:
            print(f"  VLM error frame {frame}: {e}")
            frame += interval
            continue
        elapsed = time.time() - t0
        times.append(elapsed)
        # Get embedding
        emb = get_embedding(desc)
        emb_blob = struct.pack(f"{len(emb)}d", *emb) if emb else None
        conn.execute(
            "INSERT INTO frames (file_uuid, frame, time_sec, description, embedding) VALUES (?,?,?,?,?)",
            (file_uuid, frame, frame/fps, desc, emb_blob)
        )
        conn.commit()
        total += 1
        print(f"  F{frame} ({frame/fps:.1f}s) [{elapsed:.1f}s] {desc[:80]}")
        frame += interval
    avg = sum(times)/len(times) if times else 0
    print(f"\nDone: {total} frames indexed, avg {avg:.2f}s per VLM call")

def cmd_status():
    conn = init_db()
    rows = conn.execute("""
        SELECT file_uuid, COUNT(*), MIN(frame), MAX(frame), MIN(created_at)
        FROM frames GROUP BY file_uuid
    """).fetchall()
    total = conn.execute("SELECT COUNT(*) FROM frames").fetchone()[0]
    print(f"Total indexed frames: {total}")
    for r in rows:
        print(f"  {r[0][:12]}: {r[1]} frames (F{r[2]}–F{r[3]})")

def cmd_search(query, top_k=10):
    conn = init_db()
    # Get query embedding
    q_emb = get_embedding(query)
    if not q_emb:
        # Fallback to keyword search
        print("Embedding failed, using keyword search fallback")
        rows = conn.execute(
            "SELECT file_uuid, frame, time_sec, description FROM frames WHERE description LIKE ? LIMIT ?",
            (f"%{query}%", top_k)
        ).fetchall()
        if not rows:
            print("No results")
            return
        for r in rows:
            print(f"  [{r[0][:12]}] F{r[1]} ({r[2]:.1f}s): {r[3][:100]}")
        return
    # Cosine similarity search
    rows = conn.execute("SELECT id, file_uuid, frame, time_sec, description, embedding FROM frames").fetchall()
    scored = []
    for r in rows:
        emb_blob = r[5]
        if not emb_blob or len(emb_blob) < 8:
            continue
        emb = struct.unpack(f"{len(emb_blob)//8}d", emb_blob)
        sim = cosine_similarity(q_emb, emb)
        scored.append((sim, r[1], r[2], r[3], r[4]))
    scored.sort(key=lambda x: -x[0])
    print(f"Top {min(top_k, len(scored))} results for '{query}':")
    for sim, fu, fr, ts, desc in scored[:top_k]:
        print(f"  [{sim:.3f}] [{fu[:12]}] F{fr} ({ts:.1f}s): {desc[:100]}")

def cmd_benchmark(file_uuid, count=5):
    """Benchmark VLM speed by processing N frames."""
    times = []
    for i in range(count):
        frame = i * 120
        try:
            img = get_thumbnail_proxy(file_uuid, frame)
        except:
            print(f"  Frame {frame}: not available")
            continue
        t0 = time.time()
        desc = vlm_describe(img)
        elapsed = time.time() - t0
        times.append(elapsed)
        print(f"  F{frame}: {elapsed:.2f}s — {desc[:60]}")
    if times:
        print(f"\nAvg: {sum(times)/len(times):.2f}s, Min: {min(times):.2f}s, Max: {max(times):.2f}s, Total: {sum(times):.1f}s")

# ─── Main ───
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="VLM Semantic Search for Video Keyframes")
    sub = parser.add_subparsers(dest="cmd")
    p_idx = sub.add_parser("index", help="Index a video's keyframes")
    p_idx.add_argument("file_uuid")
    p_idx.add_argument("--interval", type=int, default=FRAME_INTERVAL)
    p_idx.add_argument("--limit", type=int)

    p_search = sub.add_parser("search", help="Search indexed frames")
    p_search.add_argument("query")
    p_search.add_argument("--top-k", type=int, default=10)

    sub.add_parser("status", help="Show index status")

    p_bench = sub.add_parser("benchmark", help="Benchmark VLM speed")
    p_bench.add_argument("file_uuid")
    p_bench.add_argument("--count", type=int, default=5)

    args = parser.parse_args()
    if args.cmd == "index":
        cmd_index(args.file_uuid, args.interval, args.limit)
    elif args.cmd == "search":
        cmd_search(args.query, args.top_k)
    elif args.cmd == "status":
        cmd_status()
    elif args.cmd == "benchmark":
        cmd_benchmark(args.file_uuid, args.count)
    else:
        parser.print_help()