"""Live speech-gated MLX captions using actual causal CPU VAD decisions."""
import hashlib
import json
from pathlib import Path
import struct
import sys
import time
import mlx.core as mx
from .silero_input import GatedPCM,Silero,MODEL_SHA256,POLICY
from .stream import TextDecoder,require
from .weights import sha256


def run(args,model,frontend,profile,load_start,warmup_seconds,warmup_calls):
    from .cli import emit,memory
    import importlib.metadata
    vad=Silero(args.silero_vad)
    # Prepared before Ready, as in the ordinary CLI. Future segment resets are
    # included in the audio clock and measured separately from model windows.
    session=model.session(delay_tokens=3)
    tokenizer=TextDecoder(args.tokenizer)
    mx.synchronize()
    stream=sys.stdin.buffer if args.input=='-' else open(args.input,'rb')
    emit('caption_stream_ready',backend='mlx_metal',mlx=importlib.metadata.version('mlx'),
        weights=model.weights.provenance,weight_resident_bytes=model.weights.nbytes,
        config_sha256=sha256(args.config),frontend=frontend.provenance,tokenizer_sha256=tokenizer.table_sha256,
        source_sha256={p.name:sha256(p) for p in sorted(Path(__file__).parent.glob('*.py'))},
        silero_sha256=MODEL_SHA256,silero_runtime='onnxruntime_cpu',silero_runtime_version=vad.version,silence_policy=POLICY,
        startup='realtime18',delay_ms=240,batch=args.batch_windows,first_batch=1,dtype=args.dtype,cache_dtype=args.cache_dtype,
        math_mode=args.math_mode,fuse_projections=args.fuse_projections,projection_warmup_seconds=warmup_seconds,
        projection_warmup_calls=warmup_calls,input_mode='stdin' if args.input=='-' else ('paced_file' if args.pace_audio else 'unpaced_file'),
        load_seconds=time.monotonic()-load_start,memory=memory(),production_accuracy_approved=False)
    started=time.monotonic();source=GatedPCM(stream,vad.probability,started if args.pace_audio else None)
    ids=hashlib.sha256();emissions=0;total_work=total_reset=0.;text_bytes=0;peak_cache=0
    prior_visible=segment_visible=False;active=None;local_index=0;previous=None
    def text_event(delta,segment,**fields):
        nonlocal prior_visible,segment_visible,text_bytes
        if not delta:return
        seconds=time.monotonic()-started
        if not segment_visible and prior_visible:
            emit('caption_text_delta',segment=segment,index=emissions,seconds=seconds,text='\n',segment_separator=True)
            text_bytes+=1
        segment_visible=prior_visible=True;text_bytes+=len(delta.encode('utf8'))
        emit('caption_text_delta',segment=segment,index=emissions,seconds=seconds,text=delta,**fields)
    try:
        events=iter(source);pending=None
        while True:
            event=pending if pending is not None else next(events)
            pending=None
            if event['kind']=='start':
                s=event['segment'];require(active is None,'overlapping MLX caption session')
                active=s['id'];local_index=0;previous=None;segment_visible=False
                began=time.monotonic()
                if active:
                    session=None;mx.clear_cache();session=model.session(delay_tokens=3)
                    tokenizer=tokenizer.fresh()
                reset=time.monotonic()-began;total_reset+=reset
                require(session.position==0 and all(c.length==0 for c in session.encoder+session.decoder),'fresh MLX caches')
                emit('caption_segment_start',segment=active,start_sample=s['start_sample'],decision_sample=s['decision_sample'],
                    seconds=time.monotonic()-started,reset_seconds=reset)
            elif event['kind']=='window':
                require(active==event['segment']['id'] and local_index==event['index'],'MLX caption source clock')
                inputs=profile.prompt() if local_index==0 else [previous]
                capacity=session.batch_capacity(1 if local_index==0 else args.batch_windows,len(inputs))
                group=[event]
                while len(group)<capacity:
                    following=next(events)
                    if following['kind']!='window':
                        pending=following;break
                    require(following['segment']['id']==active and following['index']==local_index+len(group),'batch crosses segment')
                    group.append(following)
                began=time.monotonic();available=max(w['available_sample'] for w in group);ready=available/16000
                queue=max(0.,began-started-ready) if args.pace_audio or args.input=='-' else None
                first_local=local_index;first_global=emissions;last_callback=began
                def token_ready(offset,token):
                    nonlocal previous,peak_cache,emissions,local_index,last_callback
                    require(type(token) is int and 0<=token<tokenizer.model_count and token!=151645,'MLX greedy token contract')
                    require(offset==local_index-first_local and emissions==first_global+offset,'MLX batch callback clock')
                    require(session.position+session.trims*38==18+local_index,'MLX absolute position differs')
                    now=time.monotonic();interval=now-last_callback;last_callback=now
                    previous=token;ids.update(struct.pack('<I',token));peak_cache=max(peak_cache,session.cache_bytes)
                    text_event(tokenizer.push(token),active)
                    w=group[offset]
                    emit('caption_window',index=emissions,segment=active,segment_index=local_index,token_id=token,
                        seconds=time.monotonic()-started,source_dependency_sample=w['nominal_source_end_sample'],
                        available_source_sample=available,eof_drain=w['nominal_source_end_sample']>available,queue_seconds=queue,
                        work_seconds=interval,cache_bytes=session.cache_bytes,rolling_trims=session.trims,
                        batch_first_index=first_global,batch_windows=len(group),batch_offset=offset)
                    emissions+=1;local_index+=1
                session.step_batch([w['audio'] for w in group],inputs,on_token=token_ready)
                require(local_index-first_local==len(group),'MLX missing batch callback')
                total_work+=time.monotonic()-began
            elif event['kind']=='end':
                require(active==event['segment']['id'] and local_index==event['windows'],'MLX caption end clock')
                text_event(tokenizer.finish(),active,utf8_final_flush=True)
                emit('caption_segment_end',segment=active,start_sample=event['segment']['start_sample'],
                    end_sample=event['end_sample'],final_eof=event['final_eof'],unconsumed_samples=event['unconsumed_samples'],
                    windows=local_index,seconds=time.monotonic()-started)
                active=None
            elif event['kind']=='complete':break
        require(active is None and source.reader.eof,'MLX caption incomplete EOF')
        audio_seconds=source.reader.samples/16000
        emit('caption_stream_complete',status='complete',audio_samples=source.reader.samples,audio_seconds=audio_seconds,
            segments=source.segment_count,selected_samples=source.selected_samples,unconsumed_selected_samples=source.unconsumed_samples,
            emissions=emissions,text_bytes=text_bytes,model_seconds=total_work,reset_seconds=total_reset,vad_seconds=source.vad_seconds,
            read_seconds=source.read_seconds,sleep_seconds=source.sleep_seconds,elapsed_seconds=time.monotonic()-started,
            work_rtf=(total_work+total_reset+source.vad_seconds)/audio_seconds if audio_seconds else None,
            peak_cache_bytes=peak_cache,max_retained_input_samples=source.reader.max_retained_samples,
            input_sha256=source.reader.digest.hexdigest(),generated_ids_sha256=ids.hexdigest(),memory=memory(),
            all_source_audio_observed=True,measured_word_aligned_latency=False,production_accuracy_approved=False)
    finally:
        if args.input!='-':stream.close()