"""Apollo 11 mission audio search — gradio.Server backend + custom frontend.

Search the diarized transcripts (davanstrien/apollo-11-diarized) and listen
to the original tape at the matched moment (Internet Archive audio, media
fragments). See index.html for the frontend.
"""

import os
import re
from pathlib import Path

from datasets import load_dataset
from fastapi.responses import HTMLResponse, JSONResponse
from gradio import Server

DATASET = "davanstrien/apollo-11-diarized"

print("Loading segments...")
df = load_dataset(
    DATASET, "segments", split="train", token=os.environ.get("HF_TOKEN")
).to_pandas()
df["text_lower"] = df["text"].str.lower()
TAPES = sorted(df["tape"].unique().tolist())
tapes_ds = load_dataset(
    DATASET, "tapes", split="train", token=os.environ.get("HF_TOKEN")
).to_pandas()
STATS = {
    "segments": int(len(df)),
    "tapes": len(TAPES),
    "hours": int(round(tapes_ds["duration_s"].sum() / 3600)),
}
print(f"{STATS['segments']} segments across {STATS['tapes']} tapes")

app = Server()

INDEX = Path(__file__).parent / "index.html"


def fmt_time(s: float) -> str:
    h, rem = divmod(int(s), 3600)
    m, sec = divmod(rem, 60)
    return f"{h}:{m:02d}:{sec:02d}"


@app.get("/api/search")
def search(q: str = "", tape: str = "", limit: int = 150):
    hits = df
    if tape and tape in TAPES:
        hits = hits[hits["tape"] == tape]
    q = (q or "").strip().lower()
    if q:
        try:
            mask = hits["text_lower"].str.contains(q, regex=True, na=False)
        except re.error:
            mask = hits["text_lower"].str.contains(re.escape(q), na=False)
        hits = hits[mask]
    total = int(len(hits))
    rows = [
        {
            "tape": r.tape,
            "get": fmt_time(r.start),
            "start": int(r.start),
            "speaker": r.speaker or "?",
            "text": r.text,
        }
        for r in hits.head(limit).itertuples()
    ]
    return JSONResponse({"total": total, "shown": len(rows), "results": rows})


@app.get("/api/meta")
def meta():
    return JSONResponse({"stats": STATS, "tapes": TAPES})


@app.get("/", response_class=HTMLResponse)
async def homepage():
    return INDEX.read_text(encoding="utf-8")


if __name__ == "__main__":
    app.launch()