"""
gemini_batch.py - run Gemini's structured listening (gemini_audio.describe_clip) over every
recording in the workbook and cache the results.

    GEMINI_API_KEY_FILE=... python gemini_batch.py --audio-dir <dir> [--workers 4] [--retest 40]

Outputs (artifacts/features/):
    gemini_raw/<file>.json      one cached JSON per recording (resumable)
    gemini_semantic.csv         flat numeric features per file (gemini_audio.semantic_features)
    gemini_retest.csv           optional: a second, independent call for N files to measure
                                test-retest agreement of the structured fields
"""
import argparse
import json
import os
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

import pandas as pd

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import acoustic as ac        # noqa: E402
import data as bd            # noqa: E402
import gemini_audio as ga    # noqa: E402


def _one(path, client):
    y, sr = ac.load_audio(path, ac.SR)
    return ga.describe_clip(y, sr, client=client)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--audio-dir", required=True)
    ap.add_argument("--workbook", default=os.path.join(HERE, "AI_Data_Training.xlsx"))
    ap.add_argument("--out", default=os.path.join(HERE, "artifacts", "features"))
    ap.add_argument("--workers", type=int, default=4)
    ap.add_argument("--retest", type=int, default=0)
    ap.add_argument("--limit", type=int, default=None)
    a = ap.parse_args()
    if not ga.available():
        print("no Gemini API key configured (GEMINI_API_KEY / GEMINI_API_KEY_FILE); nothing to do")
        return
    raw_dir = os.path.join(a.out, "gemini_raw")
    os.makedirs(raw_dir, exist_ok=True)
    df = bd.load_workbook(a.workbook, audio_dir=a.audio_dir)
    files = list(df.loc[df["has_audio"], "sound_file"])
    on_disk = {f.lower(): f for f in os.listdir(a.audio_dir)}
    files = [on_disk[f.lower()] for f in files]
    if a.limit:
        files = files[: a.limit]
    client = ga.get_client()

    todo = [f for f in files if not os.path.exists(os.path.join(raw_dir, f + ".json"))]
    print(f"{len(files)} files, {len(todo)} to describe, workers={a.workers}", flush=True)
    t0 = time.time()
    n_err = 0
    with ThreadPoolExecutor(max_workers=a.workers) as ex:
        futs = {ex.submit(_one, os.path.join(a.audio_dir, f), client): f for f in todo}
        for k, fut in enumerate(as_completed(futs), 1):
            f = futs[fut]
            try:
                d = fut.result()
                json.dump(d, open(os.path.join(raw_dir, f + ".json"), "w"))
            except Exception as e:  # noqa
                n_err += 1
                print(f"  error {f}: {type(e).__name__}: {str(e)[:120]}", flush=True)
            if k % 50 == 0 or k == len(todo):
                print(f"  {k}/{len(todo)}  {time.time()-t0:.0f}s  errors={n_err}", flush=True)

    rows = []
    for f in files:
        p = os.path.join(raw_dir, f + ".json")
        d = json.load(open(p)) if os.path.exists(p) else None
        row = dict(file=f, **ga.semantic_features(d))
        if d is not None:
            row["gem_model"] = d.get("model")
            row["gem_dominant_sounds"] = "; ".join(d.get("dominant_sounds") or [])
            row["gem_notes"] = d.get("notes")
        rows.append(row)
    pd.DataFrame(rows).to_csv(os.path.join(a.out, "gemini_semantic.csv"), index=False)
    print("wrote gemini_semantic.csv", flush=True)

    if a.retest:
        import random
        rnd = random.Random(0)
        sub = rnd.sample(files, min(a.retest, len(files)))
        rows = []
        with ThreadPoolExecutor(max_workers=a.workers) as ex:
            futs = {ex.submit(_one, os.path.join(a.audio_dir, f), client): f for f in sub}
            for fut in as_completed(futs):
                f = futs[fut]
                try:
                    d = fut.result()
                    rows.append(dict(file=f, **ga.semantic_features(d)))
                except Exception as e:  # noqa
                    print(f"  retest error {f}: {e}", flush=True)
        pd.DataFrame(rows).to_csv(os.path.join(a.out, "gemini_retest.csv"), index=False)
        print("wrote gemini_retest.csv", flush=True)


if __name__ == "__main__":
    main()