#!/usr/bin/env python3
"""Pseudo-labeling : decode shards unlabeled avec MMS-joint, garde clips haute-confiance -> FLAC + manifest."""
import argparse, glob, os, json, subprocess
import pyarrow.parquet as pq
import soundfile as sf, numpy as np, torch
from transformers import AutoModelForCTC, AutoProcessor
SR=16000
def dec16k(raw):
    p=subprocess.run(["ffmpeg","-v","error","-i","pipe:0","-f","f32le","-ac","1","-ar",str(SR),"pipe:1"],input=raw,capture_output=True)
    return np.frombuffer(p.stdout,dtype=np.float32)
def norm(s): return " ".join(str(s).replace("|"," ").split())
def main():
    ap=argparse.ArgumentParser()
    ap.add_argument("--lang",required=True); ap.add_argument("--parquets",required=True)
    ap.add_argument("--model",default="/root/models/mmsjoint_best")
    ap.add_argument("--out_audio",required=True); ap.add_argument("--out_manifest",required=True)
    ap.add_argument("--min_conf",type=float,default=0.65); ap.add_argument("--max_clips",type=int,default=25000)
    a=ap.parse_args(); os.makedirs(a.out_audio,exist_ok=True)
    proc=AutoProcessor.from_pretrained(a.model); blank=proc.tokenizer.pad_token_id
    m=AutoModelForCTC.from_pretrained(a.model,torch_dtype=torch.bfloat16).cuda().eval()
    files=sorted(glob.glob(a.parquets)); kept=seen=0
    fo=open(a.out_manifest,"w",encoding="utf-8")
    for pfp in files:
        for b in pq.ParquetFile(pfp).iter_batches(batch_size=16):
            auds=[];metas=[]
            for r in b.to_pylist():
                au=r.get("audio"); raw=au.get("bytes") if isinstance(au,dict) else None
                if raw is None: continue
                try: wav=dec16k(raw)
                except: continue
                dur=len(wav)/SR
                if not(1.5<=dur<=30): continue
                seen+=1; auds.append(wav); metas.append((str(r.get("id") or f"{a.lang}_ul_{seen}"),dur))
            if not auds: continue
            with torch.inference_mode():
                f=proc(auds,sampling_rate=SR,return_tensors="pt",padding=True)
                f={k:(v.to("cuda",dtype=torch.bfloat16) if v.dtype==torch.float32 else v.to("cuda")) for k,v in f.items()}
                lg=m(**f).logits.float(); pr=lg.softmax(-1); mp,ids=pr.max(-1)
            txts=proc.batch_decode(ids.cpu().numpy()); idsc=ids.cpu()
            for i,(rid,dur) in enumerate(metas):
                nb=idsc[i]!=blank
                conf=float(mp[i][nb].mean()) if nb.any() else 0.0
                txt=norm(txts[i])
                if conf<a.min_conf or len(txt.split())<3: continue
                p=os.path.join(a.out_audio,rid.replace("/","_")+".flac"); sf.write(p,auds[i],SR)
                fo.write(json.dumps({"id":rid,"audio":p,"text":txt,"duration":round(dur,3)},ensure_ascii=False)+"\n"); kept+=1
            if kept>=a.max_clips: break
        if kept>=a.max_clips: break
    fo.close(); print(f"PSEUDO {a.lang}: kept {kept}/{seen}",flush=True)
if __name__=="__main__": main()