#!/usr/bin/env python3
"""istrain — voice transcription worker (VM side).

Watches the airband capture dir for new voice clips and sends each to the Mill's
whisper service (istrain-whisper stack, Dockge on the Mill). CPU lives on the Mill;
this worker only enhances (ffmpeg comms filter) and ships. Results land as:
  - a sidecar <clip>.txt next to the mp3 (same convention as radio/transcribe)
  - a durable line in ~/istrain/transcripts.jsonl (clips in /tmp are ephemeral;
    the transcript record survives reboots — keep-everything policy)

Lineage: projects/radio/transcribe/transcribe-worker.sh (marine post) — whisper.cpp
run locally there; here the model runs on the Mill and VAD replaces the prompt-echo
fixes (radio/transcribe/NOTES.md — VAD was the named "robust upgrade").
"""
import json, os, subprocess, sys, time, urllib.request, urllib.error, uuid

WATCH = os.environ.get("CLIPS_DIR", "/tmp/airband")
MILL = os.environ.get("MILL_ASR", "http://whisper:9000")
ASR = MILL + "/asr?task=transcribe&language=en&vad_filter=true&output=txt"
# env-overridable for the Mill container (writes /data/transcripts.jsonl); VM default unchanged.
OUT_LOG = os.path.expanduser(os.environ.get("TRANSCRIPTS_LOG", "~/istrain/transcripts.jsonl"))
MIN_BYTES = 5000          # below this it's a squelch blip (marine rule)
SETTLE_S = 5              # still being written if newer than this
POLL_S = 10
NO_SPEECH = "(no speech)"

def enhance(src: str, dst: str) -> bool:
    """Comms-audio cleanup: bandpass to voice, normalize. 16k mono wav."""
    r = subprocess.run(
        ["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-i", src,
         "-af", "highpass=f=250,lowpass=f=3200,speechnorm=e=12:r=0.0001:l=1",
         "-ar", "16000", "-ac", "1", dst],
        capture_output=True)
    return r.returncode == 0 and os.path.exists(dst)

def transcribe(wav: str) -> str | None:
    """POST the clip to the Mill. Returns text ('' if no speech), None on error."""
    boundary = uuid.uuid4().hex
    with open(wav, "rb") as f:
        payload = f.read()
    body = (f"--{boundary}\r\nContent-Disposition: form-data; "
            f'name="audio_file"; filename="{os.path.basename(wav)}"\r\n'
            f"Content-Type: audio/wav\r\n\r\n").encode() + payload + f"\r\n--{boundary}--\r\n".encode()
    req = urllib.request.Request(ASR, data=body, method="POST",
        headers={"Content-Type": f"multipart/form-data; boundary={boundary}"})
    try:
        with urllib.request.urlopen(req, timeout=240) as resp:
            text = resp.read().decode().strip()
    except (urllib.error.URLError, TimeoutError) as e:
        print(f"ASR error: {e}", flush=True)
        return None
    # whisper renders non-speech as runs of "..." — normalize to empty
    cleaned = " ".join(t for t in text.split() if set(t) - set(".")).strip()
    return cleaned

def main():
    print(f"transcribe-worker: watching {WATCH}, ASR={MILL}", flush=True)
    while True:
        try:
            # newest first: fresh clips (what the live panel + push want) never wait behind backlog
            names = sorted(os.listdir(WATCH),
                           key=lambda n: os.path.getmtime(os.path.join(WATCH, n))
                           if os.path.exists(os.path.join(WATCH, n)) else 0, reverse=True)
        except FileNotFoundError:
            time.sleep(POLL_S); continue
        now = time.time()
        for name in names:
            if not name.endswith(".mp3"):
                continue
            mp3 = os.path.join(WATCH, name)
            txt = mp3[:-4] + ".txt"
            if os.path.exists(txt):
                continue
            try:
                st = os.stat(mp3)
            except FileNotFoundError:
                continue
            if st.st_size < MIN_BYTES or now - st.st_mtime < SETTLE_S:
                continue
            wav = f"/tmp/transcribe-{os.getpid()}.wav"
            ok = enhance(mp3, wav)
            text = transcribe(wav if ok else mp3)
            if os.path.exists(wav):
                os.unlink(wav)
            if text is None:      # service unreachable — retry next poll, don't mark
                break
            with open(txt, "w") as f:
                f.write(text if text else NO_SPEECH)
            if text:              # only real speech goes in the durable log
                rec = {"ts": st.st_mtime, "file": name, "size": st.st_size, "text": text}
                with open(OUT_LOG, "a") as f:
                    f.write(json.dumps(rec) + "\n")
                print(f"TRANSCRIBED {name}: {text[:100]}", flush=True)
        time.sleep(POLL_S)

if __name__ == "__main__":
    main()