#!/usr/bin/env python3
"""Mac evidence runner. See the dated Mac evidence README for environment/setup.

Run each task in a separate process, serially, on an idle machine. Model caches
and the Python environment must live outside this checkout.
"""
import argparse
import hashlib
import importlib.metadata
import json
import math
import os
import platform
import re
import resource
import time
from pathlib import Path

# Reuse video-lab's existing cache; the media guard sets a separate HF_HOME.
os.environ["HF_HOME"] = str(Path.home() / ".cache/huggingface")

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("task", choices=["tts", "qwen", "asr-cpu", "asr-mlx", "granite", "lighton"])
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--repeats", type=int, default=3)
parser.add_argument("--audio-input", type=Path, help="Directory of the original Linux WAVs (defaults to this checkout)")
parser.add_argument("--image-input", type=Path, help="Existing invoice PNG (defaults to the committed fixture)")
args = parser.parse_args()
args.output.mkdir(parents=True, exist_ok=True)
ROOT = Path(__file__).resolve().parents[2]
AUDIO_DIR = args.audio_input or ROOT / "apps/marketing/public/audio/hugging-face-models-2026"
IMAGE_PATH = args.image_input or ROOT / "apps/marketing/public/images/blog/hugging-face-models-2026/sample-invoice.png"
NARRATION = "Your weekly marketing report is ready. We published three articles, updated the product page, and prepared two videos for review. Next week, we will check which work brought people to your website."
STRESS = "Teamday's Q4 report is ready. Revenue reached $1,249.50 on October 31, 2026. Review the API results at 9:30 a.m."
FILES = ["kokoro-af_heart-narration.wav", "kokoro-am_michael-narration.wav", "kokoro-af_heart-narration-noise-15db.wav"]


def digest(path):
    return hashlib.file_digest(Path(path).open("rb"), "sha256").hexdigest()


def snapshot(repo, revision):
    from huggingface_hub import snapshot_download
    path = Path(snapshot_download(repo, revision=revision))
    # Record every file used by this snapshot, including tokenizer/config assets.
    manifest = {str(p.relative_to(path)): {"bytes": p.stat().st_size, "sha256": digest(p)} for p in sorted(path.rglob("*")) if p.is_file()}
    (args.output / (repo.replace("/", "--") + "-files.json")).write_text(json.dumps(manifest, indent=2) + "\n")
    return path


def pinned(repo):
    data = json.loads((args.output / "sources.json").read_text())
    return next(m["revision"] for m in data["models"] if m["id"] == repo)


def save(data, device):
    data["environment"] = {
        "platform": platform.platform(), "python": platform.python_version(), "device": device,
        "peak_process_rss_mib": round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024**2, 2),
        "memory_note": "macOS getrusage peak RSS; not total unified-memory or Metal allocation; not minimum RAM",
        "packages": {d.metadata["Name"]: d.version for d in importlib.metadata.distributions()},
    }
    if device == "mlx-metal":
        import mlx.core as mx
        data["environment"]["mlx_device"] = str(mx.default_device())
        data["environment"]["mlx_peak_allocator_mib"] = round(mx.get_peak_memory() / 1024**2, 2)
    (args.output / f"{args.task}.json").write_text(json.dumps(data, indent=2) + "\n")
    print(f"Saved {args.task}: {len(data.get('samples', []))} observations", flush=True)


def words(text):
    return re.findall(r"[a-z0-9]+(?:'[a-z0-9]+)?", text.lower())


def errors(reference, hypothesis):
    prev = list(range(len(hypothesis) + 1))
    for i, a in enumerate(reference, 1):
        curr = [i]
        for j, b in enumerate(hypothesis, 1):
            curr.append(min(prev[j] + 1, curr[-1] + 1, prev[j-1] + (a != b)))
        prev = curr
    return prev[-1]


def audio_input(filename):
    import numpy as np
    import soundfile as sf
    from scipy.signal import resample_poly
    path = AUDIO_DIR / filename
    samples, rate = sf.read(path, dtype="float32")
    if samples.ndim > 1:
        samples = samples.mean(axis=1)
    divisor = math.gcd(rate, 16000)
    return resample_poly(samples, 16000 // divisor, rate // divisor).astype(np.float32), len(samples) / rate, digest(path)


if args.task == "tts":
    import onnxruntime as ort
    import soundfile as sf
    from huggingface_hub import hf_hub_download
    from kokoro_onnx import Kokoro
    repo, revision = "fastrtc/kokoro-onnx", "8d07950c9b6c87ce6809e9bba7bd494336217c2a"
    model = hf_hub_download(repo, "kokoro-v1.0.onnx", revision=revision)
    voices = hf_hub_download(repo, "voices-v1.0.bin", revision=revision)
    opts = ort.SessionOptions()
    opts.intra_op_num_threads, opts.inter_op_num_threads = 4, 1
    start = time.perf_counter()
    engine = Kokoro.from_session(ort.InferenceSession(model, sess_options=opts, providers=["CPUExecutionProvider"]), voices)
    load_s = time.perf_counter() - start
    engine.create("The report is ready.", voice="af_heart", speed=1.0, lang="en-us")
    rows = []
    for voice in ["af_heart", "am_michael"]:
        for label, text in [("narration", NARRATION), ("stress", STRESS)]:
            for run in range(1, args.repeats + 1):
                start = time.perf_counter()
                audio, rate = engine.create(text, voice=voice, speed=1.0, lang="en-us")
                elapsed = time.perf_counter() - start
                path = args.output / f"kokoro-{voice}-{label}-r{run}.wav"
                sf.write(path, audio, rate, subtype="PCM_16")
                rows.append(dict(file=path.name, voice=voice, text=text, run=run, seconds=len(audio)/rate, generation_s=elapsed, rtf=elapsed/(len(audio)/rate), sample_rate=rate, sha256=digest(path)))
    save(dict(model=repo, revision=revision, model_sha256=digest(model), voices_sha256=digest(voices), threads=4, speed=1.0, load_s_excluding_download=load_s, warmup="one separate sentence, af_heart", timing="phonemization and synthesis; excludes saving, loading and downloads", samples=rows), "onnx-cpu")

elif args.task == "qwen":
    import numpy as np
    import soundfile as sf
    import mlx.core as mx
    from mlx_audio.tts.utils import load_model
    mx.set_default_device(mx.gpu)
    repo = "mlx-community/Qwen3-TTS-12Hz-0.6B-CustomVoice-bf16"
    revision = pinned(repo)
    path = snapshot(repo, revision)
    start = time.perf_counter()
    model = load_model(str(path))
    mx.synchronize()
    load_s = time.perf_counter() - start
    settings = dict(speaker="Ryan", language="English", temperature=0.9, top_k=50, top_p=1.0, repetition_penalty=1.05, max_tokens=2048)
    def synth(text, seed):
        mx.random.seed(seed)
        chunks = []
        for result in model.generate_custom_voice(text=text, **settings):
            mx.eval(result.audio)
            chunks.append(np.asarray(result.audio))
        mx.synchronize()
        return np.concatenate(chunks), result.sample_rate
    synth("The report is ready.", 41)
    rows = []
    for label, text in [("narration", NARRATION), ("stress", STRESS)]:
        for run in range(1, args.repeats + 1):
            start = time.perf_counter()
            audio, rate = synth(text, 42)
            elapsed = time.perf_counter() - start
            path = args.output / f"qwen-ryan-{label}-r{run}.wav"
            sf.write(path, audio, rate, subtype="PCM_16")
            rows.append(dict(file=path.name, voice="Ryan", text=text, seed=42, run=run, seconds=len(audio)/rate, generation_s=elapsed, rtf=elapsed/(len(audio)/rate), sample_rate=rate, sha256=digest(path)))
    save(dict(model=repo, revision=revision, original_model="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", precision="bfloat16 conversion", settings=settings, load_s_excluding_download=load_s, warmup="one separate sentence, seed 41", timing="full iterator, waveform materialization and synchronization; excludes WAV saving, loading and downloads", samples=rows), "mlx-metal")

elif args.task.startswith("asr"):
    if args.task == "asr-cpu":
        from faster_whisper import WhisperModel
        repo, revision = "Systran/faster-whisper-tiny.en", "0d3d19a32d3338f10357c0889762bd8d64bbdeba"
        path = snapshot(repo, revision)
        start = time.perf_counter()
        model = WhisperModel(str(path), device="cpu", compute_type="int8", cpu_threads=4, num_workers=1)
        load_s = time.perf_counter() - start
        def transcribe(audio):
            segments, _ = model.transcribe(audio, beam_size=5, language="en", vad_filter=False, temperature=0)
            return " ".join(s.text.strip() for s in segments)
        device, settings = "ctranslate2-cpu", dict(compute_type="int8", threads=4, beam_size=5, language="en", temperature=0, vad_filter=False)
    else:
        import mlx.core as mx
        from mlx_audio.stt.utils import load_model
        mx.set_default_device(mx.gpu)
        repo = "mlx-community/whisper-large-v3-turbo-asr-fp16"
        revision = pinned(repo)
        path = snapshot(repo, revision)
        start = time.perf_counter()
        model = load_model(str(path))
        mx.synchronize()
        load_s = time.perf_counter() - start
        def transcribe(audio):
            # mlx-audio 0.6.1 explicitly rejects beam search. Record greedy
            # decoding instead of silently claiming the CPU's beam setting.
            result = model.generate(audio, language="en", temperature=0, verbose=False)
            mx.synchronize()
            return result.text
        device, settings = "mlx-metal", dict(precision="float16", beam_size=None, decoding="greedy; runtime does not implement beam search", language="en", temperature=0, vad_filter=False)
    transcribe(audio_input(FILES[0])[0])
    rows = []
    for filename in FILES:
        for run in range(1, args.repeats + 1):
            start = time.perf_counter()
            audio, duration, checksum = audio_input(filename)
            transcript = transcribe(audio)
            elapsed = time.perf_counter() - start
            count = errors(words(NARRATION), words(transcript))
            rows.append(dict(file=filename, input_sha256=checksum, run=run, transcript=transcript, reference=NARRATION, reference_words=len(words(NARRATION)), word_errors=count, wer=count/len(words(NARRATION)), transcription_s=elapsed, audio_s=duration, rtf=elapsed/duration))
    save(dict(model=repo, revision=revision, settings=settings, load_s_excluding_download=load_s, warmup="one full clean heart transcription", normalization="re.findall(r\"[a-z0-9]+(?:'[a-z0-9]+)?\", text.lower()); word Levenshtein distance / reference words", reference_limitation="intended synthetic narration, not a human-audited transcript", timing="WAV decoding, float32 polyphase resampling to 16kHz, full transcription and synchronization; excludes loading/downloads", samples=rows), device)

else:
    import torch
    from PIL import Image
    from transformers import AutoProcessor, AutoModelForImageTextToText
    torch.set_num_threads(4)
    device = "mps"
    if not torch.backends.mps.is_available():
        raise RuntimeError("This run requires Apple MPS; do not silently substitute CPU")
    repo = "ibm-granite/granite-docling-258M" if args.task == "granite" else "lightonai/LightOnOCR-2-1B"
    revision = "982fe3b40f2fa73c365bdb1bcacf6c81b7184bfe" if args.task == "granite" else pinned(repo)
    path = snapshot(repo, revision)
    image_path = IMAGE_PATH
    image = Image.open(image_path).convert("RGB")
    start = time.perf_counter()
    if args.task == "granite":
        # Pin the portable backend selected by the first successful setup.
        processor = AutoProcessor.from_pretrained(path, backend="pil")
        model_class = AutoModelForImageTextToText
    else:
        # Follow the model card's explicit classes: this checkpoint's older
        # metadata otherwise selects generic Pixtral/Mistral3 auto classes.
        from transformers import LightOnOcrProcessor, LightOnOcrForConditionalGeneration
        processor = LightOnOcrProcessor.from_pretrained(path)
        model_class = LightOnOcrForConditionalGeneration
    model = model_class.from_pretrained(path, dtype=torch.float32, attn_implementation="sdpa").eval().to(device)
    torch.mps.synchronize()
    load_s = time.perf_counter() - start
    content = [{"type": "image", "image": image}]
    if args.task == "granite":
        content.append({"type": "text", "text": "Convert this page to docling."})
    inputs = processor.apply_chat_template([{"role": "user", "content": content}], add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt").to(device)
    def infer():
        with torch.inference_mode():
            output = model.generate(**inputs, max_new_tokens=1024, do_sample=False)
        torch.mps.synchronize()
        return output[:, inputs.input_ids.shape[1]:]
    warm_start = time.perf_counter()
    infer()
    warm_s = time.perf_counter() - warm_start
    rows = []
    for run in range(1, args.repeats + 1):
        start = time.perf_counter()
        generated = infer()
        elapsed = time.perf_counter() - start
        raw = processor.batch_decode(generated, skip_special_tokens=False)[0]
        file = args.output / f"{args.task}-r{run}.txt"
        file.write_text(raw)
        rows.append(dict(run=run, file=file.name, sha256=digest(file), output=raw, generation_s=elapsed, generated_tokens=generated.shape[1], hit_token_limit=generated.shape[1]>=1024, contains_expected_values={v: v in raw for v in ["TD-2026-1031", "249.00", "49.00", "298.00", "14 November 2026"]}))
    save(dict(model=repo, revision=revision, dtype="float32", attention="sdpa", image_sha256=digest(image_path), image_size=image.size, image_processor=type(processor.image_processor).__name__, input_shapes={k: list(v.shape) for k, v in inputs.items() if hasattr(v, "shape")}, pytorch_enable_mps_fallback=os.environ.get("PYTORCH_ENABLE_MPS_FALLBACK", "unset"), max_new_tokens=1024, do_sample=False, load_s_excluding_download=load_s, warmup_generation_s=warm_s, timing="model.generate plus MPS synchronization; excludes preprocessing, decoding, saving, loading/downloads", mps_current_allocated_mib=torch.mps.current_allocated_memory()/1024**2, mps_driver_allocated_mib=torch.mps.driver_allocated_memory()/1024**2, samples=rows), "pytorch-mps")