"""Isolated Nemotron-3-Diarization service; no private audio on disk.

Run with the pinned newer Transformers in an isolated PYTHONPATH. Never upgrade
the shared ASR/TTS environment. Streaming cache and a sub-second audio buffer
are per session, in RAM; idle sessions expire, and disconnect explicitly clears.
"""
import asyncio,json,os,re,time
from pathlib import Path
import numpy as np
import torch
from aiohttp import web
from transformers import AutoModelForAudioFrameClassification,AutoProcessor

ROOT=Path(__file__).resolve().parent
MODEL=Path(os.environ.get('ROBO_DIARIZATION_MODEL','/mnt/nvme/robo-speaker-studies/diarization-model'))
DEVICE=os.environ.get('ROBO_DIARIZATION_DEVICE','cuda')
STATE={'ready':False,'sessions':{},'model':None,'processor':None,'times':[]}
LOCK=asyncio.Lock()

def local(request):
    if request.remote!='127.0.0.1' or request.host!='127.0.0.1:8796' or any(k in request.headers for k in ('CF-Ray','Forwarded','X-Forwarded-For','CF-Connecting-IP')):raise web.HTTPForbidden()

class Stream:
    def __init__(self):
        self.audio=np.zeros(0,np.float32);self.base=0;self.received=0;self.mel=0;self.cache=None;self.frames=0;self.first=True;self.seen=time.monotonic()
    @torch.inference_mode()
    def feed(self,audio):
        p=STATE['processor'];m=STATE['model'];self.audio=np.r_[self.audio,audio];self.received+=len(audio);self.seen=time.monotonic()
        start_frame=self.frames;probs=[];started=time.perf_counter()
        while True:
            start=0 if self.first else p.audio_chunk_start(self.mel)
            n=p.num_samples_first_audio_chunk if self.first else p.num_samples_per_audio_chunk
            if start+n>self.received:break
            wave=self.audio[start-self.base:start-self.base+n]
            if len(wave)!=n:raise RuntimeError('Streaming input-buffer alignment changed')
            inputs=p(wave,sampling_rate=16000,is_streaming=True,is_first_audio_chunk=self.first)
            inputs=inputs.to(m.device,dtype=m.dtype)
            output=m(**inputs,speaker_cache=self.cache);self.cache=output.speaker_cache
            values=output.logits.sigmoid()[0].float().cpu().numpy();probs.append(values);self.frames+=len(values)
            self.first=False;self.mel+=p.num_mel_frames_per_step
        next_start=0 if self.first else p.audio_chunk_start(self.mel)
        cut=max(0,next_start-self.base)
        if cut:self.audio=self.audio[cut:].copy();self.base=next_start
        elapsed=time.perf_counter()-started;STATE['times'].append(elapsed);STATE['times']=STATE['times'][-500:]
        values=np.concatenate(probs) if probs else np.zeros((0,8),np.float32)
        return {'start_frame':start_frame,'frame_seconds':.01,'probabilities':values.tolist(),
            'emitted_until_s':self.frames*.01,'received_seconds':self.received/16000,'compute_seconds':elapsed,
            'buffer_samples':len(self.audio),'input_latency_mode':'ultra_low_latency','speaker_cache_is_session_local':True}

async def load(app):
    def work():
        torch.set_num_threads(4)
        torch.backends.cuda.matmul.allow_tf32=False;torch.backends.cudnn.allow_tf32=False
        if DEVICE.startswith('cuda'):torch.cuda.set_per_process_memory_fraction(.08,0)
        p=AutoProcessor.from_pretrained(MODEL,local_files_only=True);p.set_streaming_mode('ultra_low_latency')
        m=AutoModelForAudioFrameClassification.from_pretrained(MODEL,local_files_only=True,dtype=torch.float16 if DEVICE.startswith('cuda') else torch.float32).to(DEVICE).eval()
        return p,m
    STATE['processor'],STATE['model']=await asyncio.to_thread(work);STATE['ready']=True
    async def prune():
        while True:
            await asyncio.sleep(10)
            async with LOCK:
                now=time.monotonic()
                for sid in [s for s,v in STATE['sessions'].items() if now-v.seen>60]:del STATE['sessions'][sid]
    app['pruner']=asyncio.create_task(prune())

async def cleanup(app):
    app['pruner'].cancel();await asyncio.gather(app['pruner'],return_exceptions=True);STATE['sessions'].clear()

async def state(request):
    local(request);times=STATE['times']
    return web.json_response({'ready':STATE['ready'],'model':'nvidia/Nemotron-3-Diarization','revision':'f667ed73aee57d40cc39428eb768b4fd87a0a29e',
        'device':DEVICE,'visible_gpu':os.environ.get('CUDA_VISIBLE_DEVICES'),'sessions':len(STATE['sessions']),'mode':'ultra_low_latency','audio_saved':False,
        'gpu_allocated_MiB':torch.cuda.memory_allocated()/1024**2 if DEVICE.startswith('cuda') else 0.,
        'mean_request_compute_ms':1000*float(np.mean(times)) if times else None})

async def feed(request):
    local(request)
    if not STATE['ready']:raise web.HTTPServiceUnavailable()
    sid=request.match_info['sid']
    if not re.fullmatch(r'[a-zA-Z0-9_-]{12,64}',sid):raise web.HTTPBadRequest()
    raw=await request.read()
    if not raw or len(raw)%4 or len(raw)>40960:raise web.HTTPBadRequest()
    audio=np.frombuffer(raw,'<f4')
    if not np.isfinite(audio).all():raise web.HTTPBadRequest()
    async with LOCK:
        now=time.monotonic()
        for key in [k for k,s in STATE['sessions'].items() if now-s.seen>60]:del STATE['sessions'][key]
        if sid not in STATE['sessions']:
            if len(STATE['sessions'])>=4:raise web.HTTPTooManyRequests()
            STATE['sessions'][sid]=Stream()
        result=await asyncio.to_thread(STATE['sessions'][sid].feed,audio.copy())
    return web.json_response(result)

async def close(request):
    local(request)
    async with LOCK:STATE['sessions'].pop(request.match_info['sid'],None)
    return web.json_response({'cleared':True})

def main():
    app=web.Application(client_max_size=40960);app.on_startup.append(load);app.on_cleanup.append(cleanup)
    app.router.add_get('/state',state);app.router.add_post('/session/{sid}',feed);app.router.add_delete('/session/{sid}',close)
    runtime=ROOT/'.runtime';runtime.mkdir(exist_ok=True)
    (runtime/'diarization.json').write_text(json.dumps({'pid':os.getpid(),'started':time.time(),'device':DEVICE}))
    web.run_app(app,host='127.0.0.1',port=8796,access_log=None)

if __name__=='__main__':main()