import argparse
import json
import os
import sys
import time
from pathlib import Path

# job env has a minimal PATH; whisper shells out to ffmpeg
os.environ["PATH"] = "/opt/homebrew/bin:/usr/local/bin:" + os.environ.get("PATH", "")

ap = argparse.ArgumentParser()
ap.add_argument("--input", action="append", default=[])
ap.add_argument("--outdir", required=True)
ap.add_argument("--params", default="{}")
a = ap.parse_args()
p = json.loads(a.params)

if not a.input:
    print("ERROR: needs one audio input")
    sys.exit(1)

import mlx_whisper  # noqa: E402 — import after arg errors so bad calls fail fast

t0 = time.time()
result = mlx_whisper.transcribe(
    a.input[0],
    path_or_hf_repo="mlx-community/whisper-large-v3-turbo",
    language=p.get("language") or None,
)
text = result["text"].strip()
outdir = Path(a.outdir)
stem = Path(a.input[0]).stem
(outdir / f"{stem}.txt").write_text(text + "\n")
(outdir / f"{stem}.whisper.json").write_text(json.dumps({
    "text": text,
    "language": result.get("language"),
    "segments": [{"start": s["start"], "end": s["end"], "text": s["text"].strip()}
                 for s in result["segments"]],
}, indent=2))
print(f"done in {time.time() - t0:.1f}s: {text[:150]}")