#!/usr/bin/env python3
"""Media knowledge-base benchmark: ask the same questions of a video that each
service has already ingested, record what moment it returns and/or what it
answers.

Usage: run.py --out OUT_DIR --provider NAME [--import FILE]

Providers
  echosaw     GET /v1/media/search over the account library. Needs
              ECHOSAW_API_BASE, ECHOSAW_API_KEY, ECHOSAW_MEDIA_ID (the id of the
              test video, used only to check that the right item ranked first).
  twelvelabs  search.query (moments) + analyze with pegasus (answers). Needs
              TWELVELABS_API_KEY, TWELVELABS_INDEX_ID (index holding the video).
  fireflies   No moment search on the free plan; answers come from AskFred in the
              web app and are passed with --import (JSON: {"answers": {id: text},
              "search": {id: [start_seconds, ...]}}).

Credentials are read from the environment only. Output is one JSON file per
provider with, per query, the returned moment(s) in seconds and/or the answer
text, plus wall-clock latency.
"""
import argparse
import json
import os
import time

HERE = os.path.dirname(os.path.abspath(__file__))
QUERIES = json.load(open(os.path.join(HERE, "queries.json")))["queries"]


def run_echosaw(_imported):
    import requests

    base = os.environ["ECHOSAW_API_BASE"].rstrip("/")
    h = {"X-Api-Key": os.environ["ECHOSAW_API_KEY"]}
    target = os.environ["ECHOSAW_MEDIA_ID"]
    out = {}
    library_size = None
    for q in QUERIES:
        t = time.time()
        r = requests.get(
            f"{base}/v1/media/search",
            params={"q": q["question"], "scope": "mine"},
            headers=h,
        )
        r.raise_for_status()
        d = r.json()
        lat = round(time.time() - t, 2)
        library_size = d["resultCount"]
        hits = d["results"]
        rank = next((i for i, x in enumerate(hits) if x["mediaId"] == target), None)
        t = time.time()
        a = requests.post(
            f"{base}/v1/media/ask",
            json={"question": q["question"]},
            headers=h,
        )
        a.raise_for_status()
        lat_a = round(time.time() - t, 2)
        out[q["id"]] = {
            "item_rank": None if rank is None else rank + 1,
            "moments": [hits[rank]["timestampSec"]] if rank is not None else [],
            "answer": a.json().get("answer"),
            "latency_s": lat,
            "answer_latency_s": lat_a,
        }
    return out, {
        "version": "Echosaw gamma, GET /v1/media/search?scope=mine (media-level + 60 s transcript-window vectors; one best moment per item) + POST /v1/media/ask (answers, library-wide RAG)",
        "library_items": library_size,
    }


def run_twelvelabs(_imported):
    from twelvelabs import TwelveLabs
    from twelvelabs.types import VideoContext_AssetId

    client = TwelveLabs(api_key=os.environ["TWELVELABS_API_KEY"])
    index_id = os.environ["TWELVELABS_INDEX_ID"]
    vids = list(client.indexes.videos.list(index_id))
    asset_id = None
    for v in vids:
        if getattr(v, "asset_id", None):
            asset_id = v.asset_id
    out = {}
    for q in QUERIES:
        t = time.time()
        res = list(
            client.search.query(
                index_id=index_id,
                query_text=q["question"],
                search_options=["visual", "audio"],
            )
        )
        lat_s = round(time.time() - t, 2)
        t = time.time()
        a = client.analyze(
            model_name="pegasus1.5",
            video=VideoContext_AssetId(type="asset_id", asset_id=asset_id),
            prompt=q["question"] + " Answer in one sentence.",
            temperature=0.2,
        )
        lat_a = round(time.time() - t, 2)
        out[q["id"]] = {
            "item_rank": 1 if res else None,
            "moments": [[round(x.start, 1), round(x.end, 1)] for x in res[:3]],
            "answer": a.data if isinstance(a.data, str) else json.dumps(a.data),
            "latency_s": lat_s,
            "answer_latency_s": lat_a,
        }
    return out, {
        "version": "marengo3.0 search (visual+audio) + pegasus1.5 analyze, twelvelabs SDK 1.3.4",
        "library_items": len(vids),
    }


def run_fireflies(imported):
    d = imported
    out = {}
    for q in QUERIES:
        out[q["id"]] = {
            "item_rank": 1 if d["search"].get(q["id"]) else None,
            "moments": d["search"].get(q["id"], []),
            "answer": d["answers"].get(q["id"]),
            "latency_s": None,
        }
    return out, {
        "version": d.get("version", "Fireflies web app, Free plan; AskFred answers, transcript search"),
        "library_items": d.get("library_items"),
    }


PROVIDERS = {
    "echosaw": run_echosaw,
    "twelvelabs": run_twelvelabs,
    "fireflies": run_fireflies,
}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", required=True)
    ap.add_argument("--provider", required=True, choices=sorted(PROVIDERS))
    ap.add_argument("--import", dest="imported")
    a = ap.parse_args()
    imported = json.load(open(a.imported)) if a.imported else None
    results, meta = PROVIDERS[a.provider](imported)
    os.makedirs(a.out, exist_ok=True)
    with open(os.path.join(a.out, f"{a.provider}.json"), "w") as f:
        json.dump({"provider": a.provider, **meta, "queries": results}, f, indent=1)
    print(a.provider, "done")


if __name__ == "__main__":
    main()