#!/usr/bin/env python3
"""Reproduce the HF guide samples; dependencies and commands are in the evidence README."""
import argparse
import hashlib
import importlib.metadata
import json
import platform
import re
import resource
import time
from pathlib import Path

parser = argparse.ArgumentParser()
parser.add_argument("task", choices=["tts", "asr", "ocr"])
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
args.output.mkdir(parents=True, exist_ok=True)

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."


def digest(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def save(name, data):
    peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    data["environment"] = {
        "platform": platform.platform(),
        "processor": platform.processor(),
        "device": "cpu",
        "threads": 4,
        "peak_process_rss_mib": round(peak / (1024 if platform.system() == "Linux" else 1024 ** 2), 1),
        "packages": {d.metadata["Name"]: d.version for d in importlib.metadata.distributions()},
    }
    (args.output / name).write_text(json.dumps(data, indent=2) + "\n")
    print(json.dumps({k: v for k, v in data.items() if k != "environment"}), flush=True)


if args.task == "tts":
    import numpy as np
    import onnxruntime as ort
    import soundfile as sf
    from huggingface_hub import hf_hub_download
    from kokoro_onnx import Kokoro

    repo = "fastrtc/kokoro-onnx"
    revision = "8d07950c9b6c87ce6809e9bba7bd494336217c2a"
    model = Path(hf_hub_download(repo, "kokoro-v1.0.onnx", revision=revision))
    voices = Path(hf_hub_download(repo, "voices-v1.0.bin", revision=revision))
    options = ort.SessionOptions()
    options.intra_op_num_threads = 4
    options.inter_op_num_threads = 1
    start = time.perf_counter()
    session = ort.InferenceSession(str(model), sess_options=options, providers=["CPUExecutionProvider"])
    engine = Kokoro.from_session(session, str(voices))
    load_s = time.perf_counter() - start
    # Separate warm-up; timings below include phonemization and synthesis, not loading.
    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)]:
            start = time.perf_counter()
            samples, rate = engine.create(text, voice=voice, speed=1.0, lang="en-us")
            elapsed = time.perf_counter() - start
            path = args.output / f"kokoro-{voice}-{label}.wav"
            sf.write(path, samples, rate, subtype="PCM_16")
            rows.append({"file": path.name, "voice": voice, "text": text, "seconds": round(len(samples) / rate, 3), "generation_s": round(elapsed, 3), "rtf": round(elapsed / (len(samples) / rate), 4), "sample_rate": rate, "sha256": digest(path)})
    clean, rate = sf.read(args.output / "kokoro-af_heart-narration.wav")
    noise = np.random.default_rng(42).normal(size=clean.shape)
    noise *= np.sqrt(np.mean(clean ** 2) / (np.mean(noise ** 2) * 10 ** (15 / 10)))
    noisy = args.output / "kokoro-af_heart-narration-noise-15db.wav"
    sf.write(noisy, np.clip(clean + noise, -1, 1), rate, subtype="PCM_16")
    save("tts.json", {"model": repo, "revision": revision, "model_sha256": digest(model), "voices_sha256": digest(voices), "load_s_excluding_download": round(load_s, 3), "warmup": "one sentence, af_heart", "speed": 1.0, "samples": rows, "noise": {"file": noisy.name, "snr_db": 15, "seed": 42, "sha256": digest(noisy)}})

elif args.task == "asr":
    import math
    import numpy as np
    import soundfile as sf
    from faster_whisper import WhisperModel
    from huggingface_hub import snapshot_download
    from scipy.signal import resample_poly

    def audio_input(filename):
        samples, rate = sf.read(args.output / filename, 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)

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

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

    rows = []
    for repo, revision in [("Systran/faster-whisper-tiny.en", "0d3d19a32d3338f10357c0889762bd8d64bbdeba"), ("Systran/faster-whisper-base.en", "3d3d5dee26484f91867d81cb899cfcf72b96be6c"), ("dropbox-dash/faster-whisper-large-v3-turbo", "0a363e9161cbc7ed1431c9597a8ceaf0c4f78fcf")]:
        path = snapshot_download(repo, revision=revision)
        engine = WhisperModel(path, device="cpu", compute_type="int8", cpu_threads=4, num_workers=1)
        # Warm up the decoder before timing.
        list(engine.transcribe(audio_input("kokoro-af_heart-narration.wav"), beam_size=5, language="en", vad_filter=False)[0])
        for filename in ["kokoro-af_heart-narration.wav", "kokoro-am_michael-narration.wav", "kokoro-af_heart-narration-noise-15db.wav"]:
            start = time.perf_counter()
            segments, info = engine.transcribe(audio_input(filename), beam_size=5, language="en", vad_filter=False, temperature=0)
            transcript = " ".join(s.text.strip() for s in segments)
            elapsed = time.perf_counter() - start
            count = errors(words(NARRATION), words(transcript))
            rows.append({"model": repo, "revision": revision, "file": filename, "transcript": transcript, "reference": NARRATION, "reference_words": len(words(NARRATION)), "word_errors": count, "wer": round(count / len(words(NARRATION)), 4), "transcription_s": round(elapsed, 3), "audio_s": round(info.duration, 3), "rtf": round(elapsed / info.duration, 4)})
        del engine
    save("asr.json", {"compute_type": "int8", "beam_size": 5, "temperature": 0, "language": "en", "vad_filter": False, "audio_preprocessing": "soundfile float32 decode; scipy.signal.resample_poly from source rate to 16000 Hz; included in timing", "normalization": "lowercase, strip punctuation except internal apostrophes; Levenshtein word distance", "samples": rows})

elif args.task == "ocr":
    import torch
    from PIL import Image, ImageDraw, ImageFont
    from transformers import AutoProcessor, AutoModelForImageTextToText

    torch.set_num_threads(4)
    font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 26)
    image = Image.new("RGB", (1000, 720), "white")
    draw = ImageDraw.Draw(image)
    lines = ["SAMPLE INVOICE", "Example Company", "Invoice: TD-2026-1031", "Date: 31 October 2026", "Description                         Amount", "Content review                      USD 249.00", "Voiceover editing                   USD 49.00", "Total                               USD 298.00", "Payment due: 14 November 2026"]
    for i, line in enumerate(lines):
        draw.text((50, 45 + i * 66), line, font=font, fill="black")
    image_path = args.output / "sample-invoice.png"
    image.save(image_path)
    repo = "ibm-granite/granite-docling-258M"
    revision = "982fe3b40f2fa73c365bdb1bcacf6c81b7184bfe"
    processor = AutoProcessor.from_pretrained(repo, revision=revision)
    model = AutoModelForImageTextToText.from_pretrained(repo, revision=revision, dtype=torch.float32, attn_implementation="sdpa").eval()
    messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "Convert this page to docling."}]}]
    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
    inputs = processor(text=prompt, images=[image], return_tensors="pt")
    start = time.perf_counter()
    with torch.inference_mode():
        output = model.generate(**inputs, max_new_tokens=1024, do_sample=False)
    elapsed = time.perf_counter() - start
    generated = output[:, inputs.input_ids.shape[1]:]
    text = processor.batch_decode(generated, skip_special_tokens=False)[0]
    (args.output / "sample-invoice-doctags.txt").write_text(text)
    save("ocr.json", {"model": repo, "revision": revision, "image": image_path.name, "image_sha256": digest(image_path), "reference_lines": lines, "prompt": "Convert this page to docling.", "dtype": "float32", "max_new_tokens": 1024, "generated_tokens": generated.shape[1], "hit_token_limit": generated.shape[1] >= 1024, "generation_s": round(elapsed, 3), "warmup": "none; single first inference, excludes loading and downloads", "output": text, "contains_expected_values": {v: v in text for v in ["TD-2026-1031", "249.00", "49.00", "298.00", "14 November 2026"]}})