#!/usr/bin/env python3
"""Verify live audio indexing with a run-owned video containing clear speech."""

from __future__ import annotations

import argparse
import atexit
import hashlib
import json
import time
from pathlib import Path
from typing import Any

import httpx


def call(client: httpx.Client, method: str, path: str, **kwargs: Any) -> httpx.Response:
    response = client.request(method, path, **kwargs)
    if response.is_error:
        raise RuntimeError(f"{method} {path} -> {response.status_code}: {response.text[:1000]}")
    return response


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while block := handle.read(1024 * 1024):
            digest.update(block)
    return digest.hexdigest()


def find_video(client: httpx.Client, source_id: str, filename: str) -> dict[str, Any] | None:
    videos = call(client, "GET", "/api/v1/videos?limit=500").json()
    return next(
        (
            item
            for item in videos
            if item.get("source_id") == source_id
            and item.get("source_key", "").rsplit("/", 1)[-1] == filename
        ),
        None,
    )


def upload_fixture(
    client: httpx.Client,
    fixture: Path,
    state: dict[str, Any],
) -> tuple[dict[str, Any], list[float]]:
    filename = f"{state['run_id']}-{fixture.name}"
    uploads = call(client, "GET", "/api/v1/uploads?limit=500").json()
    upload = next((item for item in uploads if item.get("filename") == filename), None)
    if upload is None or upload.get("status") in {"failed", "cancelled"}:
        upload = call(
            client,
            "POST",
            "/api/v1/uploads",
            json={
                "source_id": state["source_id"],
                "relative_path": "ingest",
                "filename": filename,
                "title": "音频 GPU 验收",
                "collection_id": state["collection_id"],
                "tag_ids": [state["tag_id"]],
                "size_bytes": fixture.stat().st_size,
                "sha256": sha256(fixture),
                "conflict": "skip",
            },
        ).json()

    latencies: list[float] = []
    if upload.get("status") != "completed":
        hidden = find_video(client, state["source_id"], filename)
        if hidden is not None:
            raise AssertionError("unfinished speech upload leaked into /videos")
        chunk_size = int(upload["chunk_size"])
        received = {int(index) for index in upload.get("received_chunks", upload.get("received", []))}
        with fixture.open("rb") as handle:
            for index in range(int(upload["total_chunks"])):
                data = handle.read(chunk_size)
                if index in received:
                    continue
                started = time.perf_counter()
                call(
                    client,
                    "PUT",
                    f"/api/v1/uploads/{upload['id']}/chunks/{index}",
                    content=data,
                    headers={"X-Chunk-SHA256": hashlib.sha256(data).hexdigest()},
                )
                latencies.append(round((time.perf_counter() - started) * 1000, 1))
        upload_id = upload["id"]
        upload = call(client, "POST", f"/api/v1/uploads/{upload_id}/complete").json()
    else:
        upload_id = upload["id"]

    deadline = time.monotonic() + 240
    while upload.get("status") not in {"completed", "failed", "cancelled"}:
        if time.monotonic() >= deadline:
            raise TimeoutError(f"speech upload did not finish: {upload.get('status')}")
        time.sleep(1)
        uploads = call(client, "GET", "/api/v1/uploads?limit=500").json()
        upload = next(item for item in uploads if item["id"] == upload_id)
    if upload["status"] != "completed":
        raise AssertionError(f"speech upload failed: {upload.get('message')}")
    return upload, latencies


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-url", required=True)
    auth = parser.add_mutually_exclusive_group(required=True)
    auth.add_argument("--token-file", type=Path)
    auth.add_argument("--password")
    parser.add_argument("--run-dir", type=Path, required=True)
    parser.add_argument("--fixture", default="e2e-speech.mp4")
    parser.add_argument(
        "--query",
        action="append",
        dest="queries",
        help="Required audio-search keyword; repeat to verify multiple words (default: Americans)",
    )
    parser.add_argument("--timeout", type=int, default=1200)
    args = parser.parse_args()

    state_path = args.run_dir / "state.json"
    state = json.loads(state_path.read_text())
    fixture = args.run_dir / args.fixture
    if not fixture.is_file():
        raise FileNotFoundError(fixture)
    headers = {}
    if args.token_file is not None:
        headers["Authorization"] = f"Bearer {args.token_file.read_text().strip()}"
    client = httpx.Client(
        base_url=args.base_url.rstrip("/"),
        headers=headers,
        timeout=httpx.Timeout(60, connect=10),
    )
    if args.password is not None:
        login = call(
            client,
            "POST",
            "/api/v1/auth/login",
            json={"password": args.password, "remember_device": False},
        ).json()
        csrf = str(login.get("csrf_token") or "")
        if not csrf:
            raise AssertionError("login response omitted CSRF token")
        client.headers["X-CSRF-Token"] = csrf

        def logout() -> None:
            try:
                client.post("/api/v1/auth/logout")
            except Exception:
                pass

        atexit.register(logout)

    # The repository's default live fixture contains the JFK inauguration line
    # "And so, my fellow Americans...".  Keep the default assertion tied to
    # the actual fixture; callers using another recording can repeat --query.
    required_queries = [value.strip() for value in (args.queries or ["Americans"]) if value.strip()]
    if not required_queries:
        raise ValueError("at least one non-empty --query is required")

    initial_models = call(client, "GET", "/api/v1/models").json()
    print(
        json.dumps(
            {
                "audio_model": (initial_models.get("manifest") or {})
                .get("components", {})
                .get("audio"),
                "audio_health": (initial_models.get("component_health") or {}).get("audio"),
                "audio_accelerator": (initial_models.get("accelerator") or {})
                .get("components", {})
                .get("audio"),
            },
            ensure_ascii=False,
        ),
        flush=True,
    )

    upload, chunk_latencies = upload_fixture(client, fixture, state)
    filename = upload["filename"]
    deadline = time.monotonic() + args.timeout
    observed_devices: list[str] = []
    latest_message = "等待视频入库"
    positives: dict[str, dict[str, Any]] = {}
    video: dict[str, Any] | None = None
    last_report_at = 0.0
    while time.monotonic() < deadline:
        video = find_video(client, state["source_id"], filename)
        models = call(client, "GET", "/api/v1/models").json()
        audio = models.get("accelerator", {}).get("components", {}).get("audio", {})
        device = audio.get("actual_device") or audio.get("device")
        if device and device not in observed_devices:
            observed_devices.append(device)
        if video is not None:
            for query in required_queries:
                if query in positives:
                    continue
                result = call(
                    client,
                    "POST",
                    "/api/v1/search",
                    json={"text": query, "recognition_types": ["audio"]},
                ).json()
                positive = next(
                    (item for item in result.get("items", []) if item.get("video_id") == video["id"]),
                    None,
                )
                if positive is not None:
                    positives[query] = positive
            if len(positives) == len(required_queries):
                break
        jobs = call(client, "GET", "/api/v1/jobs?page=1&page_size=10").json()["items"]
        audio_jobs = [item for item in jobs if item.get("kind") == "transcribe_audio"]
        if audio_jobs:
            latest_message = audio_jobs[0].get("message") or audio_jobs[0].get("status", "")
            if audio_jobs[0].get("status") == "failed":
                raise AssertionError(f"audio job failed: {audio_jobs[0].get('error')}")
        now = time.monotonic()
        if now - last_report_at >= 10:
            print(
                json.dumps(
                    {
                        "waiting": latest_message,
                        "video_found": video is not None,
                        "observed_devices": observed_devices,
                    },
                    ensure_ascii=False,
                ),
                flush=True,
            )
            last_report_at = now
        time.sleep(2)
    if video is None:
        raise AssertionError("speech video never appeared in /videos")
    missing_queries = [query for query in required_queries if query not in positives]
    if missing_queries:
        raise TimeoutError(
            "audio search never matched required keywords "
            f"{missing_queries}; matched={sorted(positives)}; latest={latest_message}"
        )
    for query, positive in positives.items():
        if positive.get("segment_start_ms") is None or positive.get("segment_end_ms") is None:
            raise AssertionError(f"audio match omitted time segment for {query!r}: {positive}")
        match_details = positive.get("match_details", [])
        audio_text = " ".join(
            str(detail.get("text") or "")
            for detail in match_details
            if detail.get("type") == "audio"
        ).strip()
        if not audio_text:
            raise AssertionError(f"audio match omitted recognized text for {query!r}: {positive}")
        if len([character for character in audio_text if character.isalnum()]) < 8:
            raise AssertionError(
                f"audio transcript is implausibly short for {query!r}: {audio_text!r}"
            )
    negative = call(
        client,
        "POST",
        "/api/v1/search",
        json={"text": "pineapple submarine 9834721", "recognition_types": ["audio"]},
    ).json()
    if any(item.get("video_id") == video["id"] for item in negative.get("items", [])):
        raise AssertionError("non-matching audio search unexpectedly returned the speech video")

    jobs = call(client, "GET", "/api/v1/jobs?page=1&page_size=10").json()["items"]
    latest_audio = next((item for item in jobs if item.get("kind") == "transcribe_audio"), None)
    report = {
        "upload_id": upload["id"],
        "video_id": video["id"],
        "chunk_latency_ms": chunk_latencies,
        "observed_audio_devices": observed_devices,
        "audio_job": latest_audio,
        "matches": {
            query: {
                "segment_start_ms": positive.get("segment_start_ms"),
                "segment_end_ms": positive.get("segment_end_ms"),
                "details": positive.get("match_details", []),
            }
            for query, positive in positives.items()
        },
        "negative_matches": len(negative.get("items", [])),
    }
    if not any(str(device).upper().startswith("GPU") for device in observed_devices):
        raise AssertionError(f"GPU was not observed during speech indexing: {observed_devices}")
    state[f"speech_upload_id_{fixture.stem}"] = upload["id"]
    state[f"speech_video_id_{fixture.stem}"] = video["id"]
    state_path.write_text(json.dumps(state, ensure_ascii=False, indent=2) + "\n")
    (args.run_dir / "speech-report.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n"
    )
    print(json.dumps(report, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()