"""Portable 560ms explicit-cache EN reference runtime for bundled LiteRT graphs."""
from pathlib import Path
import json,time,hashlib
import numpy as np

class Graph:
    def __init__(self,path,threads=4):
        from ai_edge_litert.interpreter import Interpreter
        self.interpreter=Interpreter(model_path=str(path),num_threads=threads)
        self.runner=self.interpreter.get_signature_runner('serving_default')
    def __call__(self,*args):
        result=self.runner(**{f'args_{i}':a for i,a in enumerate(args)})
        outputs=tuple(result[k] for k in sorted(result,key=lambda k:int(k.rsplit('_',1)[1])))
        for out in outputs:
            if not np.isfinite(out).all():raise ValueError('Nonfinite graph output')
        return outputs

class Runtime:
    def __init__(self,bundle,threads=4,allow_experimental=False,validate_manifest=True):
        bundle=Path(bundle)
        if validate_manifest:
            manifest=json.loads((bundle/'nemotron-manifest.json').read_text())
            if manifest['adapter_id']!='nemotron-en-hf560-v1' or manifest['blank_id']!=1024:raise ValueError('Wrong model adapter')
            if manifest.get('evaluation_only') and not allow_experimental:raise ValueError('Experimental bundle: inspect parity report and explicitly allow evaluation')
            for entry in manifest['files']:
                path=bundle/entry['path']
                if path.parent!=bundle or not path.is_file():raise ValueError('Invalid or missing bundle file')
                if path.stat().st_size!=entry['bytes'] or hashlib.sha256(path.read_bytes()).hexdigest()!=entry['sha256']:raise ValueError('Bundle checksum mismatch: '+path.name)
        self.graphs={p.stem:Graph(p,threads) for p in Path(bundle).glob('*.tflite')}
        required={'frontend_initial','frontend','subsampling_initial','subsampling','projection','predictor','joint'}|{f'block_{i:02}' for i in range(24)}
        if set(self.graphs)!=required:raise ValueError('Incomplete or unexpected EN component inventory')
    def transcribe(self,pcm):
        if pcm.ndim!=1 or not np.isfinite(pcm).all():raise ValueError('Expected finite mono float32 PCM at 16kHz')
        z=lambda *s:np.zeros(s,np.float32)
        states=[z(1,1,1,131),z(1,256,1,68),z(1,256,1,36)]
        blocks=[[z(1,8,70,128),z(1,8,70,128),z(1,1024,8)] for _ in range(24)]
        g=self.graphs;tokens=[];emissions=[];timings={};step=0;offset=0
        def run(name,*args):
            begun=time.monotonic();out=g[name](*args);timings[name]=timings.get(name,0)+time.monotonic()-begun;return out
        pred,h,c=run('predictor',np.array([[1024]],np.int64),z(2,1,640),z(2,1,640))
        while offset<len(pcm)//160:
            frames=49 if step==0 else 56;length=frames*160+353;begin=offset*160-257
            wave=z(1,length);a=max(0,begin);b=min(len(pcm),begin+length)
            if b>a:wave[0,a-begin:b-begin]=pcm[a:b]
            valid=np.array([min(frames,max(0,len(pcm)//160-offset))],np.int32)
            samples=np.array([min(length-1,max(0,len(pcm)-begin-1))],np.int32)
            mel,=run('frontend_initial' if step==0 else 'frontend',wave,valid,samples)
            x,*states=run('subsampling_initial' if step==0 else 'subsampling',mel,*states)
            for i in range(24):x,*blocks[i]=run(f'block_{i:02}',x,*blocks[i],np.array([min(step*7,70)],np.int32))
            projected,=run('projection',x)
            for f in range(7):
                for symbol in range(10):
                    logits,=run('joint',projected[:,f:f+1],pred);token=int(logits.reshape(-1).argmax())
                    if token==1024:break
                    tokens.append(token);emissions.append((step*7+f)*.08)
                    if len(tokens)>100000:raise ValueError('Token limit')
                    pred,h,c=run('predictor',np.array([[token]],np.int64),h,c)
                else:raise ValueError('No blank within 10 evaluations on one encoder frame')
            offset+=frames;step+=1
        return dict(tokens=tokens,emission_seconds=emissions,compute_seconds=sum(timings.values()),timings=timings,steps=step)


def decode(tokens,vocabulary):
    text=''.join(vocabulary['pieces'][t] for t in tokens if t not in vocabulary['special_ids']).replace('▁',' ')
    return text[1:] if text.startswith(' ') else text


if __name__=='__main__':
    import argparse,math
    import soundfile as sf
    from scipy.signal import resample_poly
    parser=argparse.ArgumentParser();parser.add_argument('bundle',type=Path);parser.add_argument('audio',type=Path);parser.add_argument('--threads',type=int,default=4);parser.add_argument('--allow-experimental',action='store_true');args=parser.parse_args()
    wave,sr=sf.read(args.audio,dtype='float32',always_2d=True);divisor=math.gcd(sr,16000)
    pcm=resample_poly(wave.mean(axis=1),16000//divisor,sr//divisor).astype('float32')
    runtime=Runtime(args.bundle,args.threads,args.allow_experimental)
    result=runtime.transcribe(pcm);vocabulary=json.loads((args.bundle/'vocabulary.json').read_text())
    result['text']=decode(result['tokens'],vocabulary)
    print(json.dumps(result,ensure_ascii=False,indent=2))