#!/usr/bin/env python3
"""Generate draft content labels for exported WAVs with local Whisper ASR.

These labels are not native/Radium symbols. They are content-derived draft
transcripts intended to make voice-like sounds easier to triage in the wiki.
"""

from __future__ import annotations

import argparse
import csv
import re
import sys
from pathlib import Path


def clean_text(value: str) -> str:
    value = re.sub(r"\s+", " ", value).strip()
    return value.strip(" \"'`")


def draft_slug(value: str, max_words: int = 8) -> str:
    words = re.findall(r"[A-Za-z0-9]+", value.upper())
    return "_".join(words[:max_words]) or "ASR_DRAFT"


def load_existing(path: Path) -> dict[str, dict[str, str]]:
    if not path.exists():
        return {}
    with path.open(newline="") as f:
        return {row["sound_name"]: row for row in csv.DictReader(f, delimiter="\t")}


def passes_duration(row: dict[str, str], min_duration: float, max_duration: float, channels: set[int]) -> bool:
    duration = float(row["duration_seconds"])
    channel_count = int(row["channels"])
    return min_duration <= duration <= max_duration and channel_count in channels


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--manifest", type=Path, default=Path("analysis/radium-image/sound_export_manifest.tsv"))
    parser.add_argument("--out", type=Path, default=Path("analysis/radium-image/asr_draft_sound_labels.tsv"))
    parser.add_argument("--model", default="tiny.en")
    parser.add_argument("--model-dir", type=Path, default=Path("/Users/jordan/.cache/whisper"))
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--language", default="en")
    parser.add_argument("--min-duration", type=float, default=0.8)
    parser.add_argument("--max-duration", type=float, default=6.0)
    parser.add_argument("--channels", default="1,2", help="Comma-separated channel counts to include.")
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--resume", action="store_true")
    parser.add_argument("--accept-logprob", type=float, default=-0.85)
    parser.add_argument("--accept-no-speech", type=float, default=0.35)
    parser.add_argument("--accept-compression", type=float, default=2.0)
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.0,
        help="Whisper decode temperature. A single value disables the slow fallback temperature schedule.",
    )
    args = parser.parse_args()

    try:
        import whisper
    except ImportError as exc:
        raise SystemExit(
            "Python module 'whisper' is not importable. Run with the Homebrew whisper Python, "
            "for example /opt/homebrew/Cellar/openai-whisper/20250625_4/libexec/bin/python."
        ) from exc

    channels = {int(part) for part in args.channels.split(",") if part.strip()}
    manifest_rows = list(csv.DictReader(args.manifest.open(newline=""), delimiter="\t"))
    candidates = [row for row in manifest_rows if passes_duration(row, args.min_duration, args.max_duration, channels)]
    if args.limit:
        candidates = candidates[: args.limit]

    existing = load_existing(args.out) if args.resume else {}
    output_rows = [existing[name] for name in sorted(existing)] if existing else []
    done = set(existing)

    args.out.parent.mkdir(parents=True, exist_ok=True)
    fields = [
        "sound_name",
        "wav_path",
        "section8_index",
        "object_id",
        "duration_seconds",
        "channels",
        "transcript",
        "draft_label",
        "accepted",
        "avg_logprob",
        "no_speech_prob",
        "compression_ratio",
        "segment_count",
        "model",
    ]

    model = whisper.load_model(args.model, download_root=str(args.model_dir), device=args.device)
    with args.out.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fields, delimiter="\t")
        writer.writeheader()
        for row in output_rows:
            writer.writerow(row)
        f.flush()

        processed = len(done)
        for row in candidates:
            if row["sound_name"] in done:
                continue
            result = model.transcribe(
                row["wav_path"],
                language=args.language,
                fp16=False,
                verbose=False,
                condition_on_previous_text=False,
                temperature=args.temperature,
            )
            segments = result.get("segments") or []
            transcript = clean_text(str(result.get("text") or ""))
            if segments:
                total_duration = sum(max(0.01, float(seg["end"]) - float(seg["start"])) for seg in segments)
                avg_logprob = sum(
                    float(seg.get("avg_logprob", -99.0)) * max(0.01, float(seg["end"]) - float(seg["start"]))
                    for seg in segments
                ) / total_duration
                no_speech_prob = max(float(seg.get("no_speech_prob", 1.0)) for seg in segments)
                compression_ratio = max(float(seg.get("compression_ratio", 99.0)) for seg in segments)
            else:
                avg_logprob = -99.0
                no_speech_prob = 1.0
                compression_ratio = 99.0

            accepted = (
                bool(transcript)
                and avg_logprob >= args.accept_logprob
                and no_speech_prob <= args.accept_no_speech
                and compression_ratio <= args.accept_compression
                and not re.fullmatch(r"[*\\W_]+", transcript)
            )
            out_row = {
                "sound_name": row["sound_name"],
                "wav_path": row["wav_path"],
                "section8_index": row["section8_index"],
                "object_id": row["primary_object_local_id"],
                "duration_seconds": row["duration_seconds"],
                "channels": row["channels"],
                "transcript": transcript,
                "draft_label": f"ASR_{draft_slug(transcript)}" if accepted else "",
                "accepted": "yes" if accepted else "no",
                "avg_logprob": f"{avg_logprob:.6f}",
                "no_speech_prob": f"{no_speech_prob:.6f}",
                "compression_ratio": f"{compression_ratio:.6f}",
                "segment_count": len(segments),
                "model": args.model,
            }
            writer.writerow(out_row)
            f.flush()
            processed += 1
            if processed % 25 == 0:
                print(f"processed {processed}/{len(candidates)}", file=sys.stderr)

    print(f"wrote {args.out} rows={sum(1 for _ in args.out.open()) - 1} candidates={len(candidates)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
