"""Bounded, tiled Real-ESRGAN video inference. No CPU fallback or frame interpolation."""
import argparse, json, subprocess, sys, os, hashlib
from fractions import Fraction
from pathlib import Path
import numpy as np
import torch
from PIL import Image
sys.path.insert(0, str(Path(__file__).parent / 'vendor'))
from srvgg_arch import SRVGGNetCompact

def probe(file):
    return json.loads(subprocess.check_output(['ffprobe','-v','error','-protocol_whitelist','file,pipe','-show_streams','-show_format','-of','json',file], timeout=30))

def main():
    p=argparse.ArgumentParser()
    for name in ['input','output','model','resolution']: p.add_argument('--'+name, required=True)
    p.add_argument('--tile', type=int, default=192)
    a=p.parse_args()
    if not torch.cuda.is_available(): raise RuntimeError('CUDA unavailable; refusing CPU fallback')
    torch.set_num_threads(4)
    info=probe(a.input); v=next(s for s in info['streams'] if s['codec_type']=='video')
    if not any(name in info['format']['format_name'].split(',') for name in ['mov','mp4','matroska','webm','avi','flv','mpegts']): raise ValueError('Unsupported video container')
    rotation=next((int(side['rotation'])%360 for side in v.get('side_data_list',[]) if 'rotation' in side),0)
    if rotation%90: raise ValueError('Unsupported rotation angle')
    w,h=int(v['width']),int(v['height'])
    if rotation in [90,270]: w,h=h,w
    fps=Fraction(v['avg_frame_rate'])
    duration=float(info['format']['duration'])
    if not 0<duration<=120 or not 0<float(fps)<=60 or w*h>3840*2160: raise ValueError('Input exceeds 120 s / 60 fps / 4K limits')
    # Reject variable frame rate rather than silently altering timing/audio sync.
    frame_times=subprocess.check_output(['ffprobe','-v','error','-protocol_whitelist','file,pipe','-select_streams','v:0','-show_entries','frame=best_effort_timestamp_time','-of','csv=p=0',a.input],timeout=60).decode().splitlines()
    timestamps=[float(s.split(',')[0]) for s in frame_times if s.split(',')[0].strip()]
    if any(abs((b-c)-float(1/fps))>max(.002,float(1/fps)*.05) for c,b in zip(timestamps,timestamps[1:])): raise ValueError('Variable frame rate video must be converted to constant frame rate first')
    long,short={'1080P':(1920,1080),'2K':(2560,1440),'4K':(3840,2160)}[a.resolution]
    scale=min((long if w>=h else short)/w,(short if w>=h else long)/h)
    if not 1<scale<=4: raise ValueError('Target must exceed source resolution by no more than 4x')
    ow,oh=int(w*scale)//2*2,int(h*scale)//2*2
    model=SRVGGNetCompact(num_in_ch=3,num_out_ch=3,num_feat=64,num_conv=32,upscale=4,act_type='prelu')
    if hashlib.sha256(Path(a.model).read_bytes()).hexdigest() != '8dc7edb9ac80ccdc30c3a5dca6616509367f05fbc184ad95b731f05bece96292': raise RuntimeError('Model checksum mismatch')
    weights=torch.load(a.model,map_location='cpu',weights_only=True)
    model.load_state_dict(weights.get('params_ema',weights.get('params',weights)),strict=True)
    model.eval().to('cuda')
    temp=a.output+'.silent.mp4'
    decode=subprocess.Popen(['ffmpeg','-v','error','-filter_threads','1','-filter_complex_threads','1','-threads','2','-protocol_whitelist','file,pipe','-i',a.input,'-map','0:v:0','-f','rawvideo','-pix_fmt','rgb24','-vsync','0','pipe:1'],stdout=subprocess.PIPE)
    encode=subprocess.Popen(['ffmpeg','-v','error','-filter_threads','1','-filter_complex_threads','1','-y','-f','rawvideo','-pix_fmt','rgb24','-s',f'{ow}x{oh}','-r',str(fps),'-i','pipe:0','-an','-c:v','libx264','-threads','4','-preset','fast','-crf','18','-pix_fmt','yuv420p',temp],stdin=subprocess.PIPE)
    count=0
    try:
        with torch.inference_mode():
            while True:
                raw=decode.stdout.read(w*h*3)
                if not raw: break
                if len(raw)!=w*h*3: raise RuntimeError('Truncated video frame')
                source=np.frombuffer(raw,dtype=np.uint8).reshape(h,w,3)
                # CPU output buffer, bounded GPU tiles with enough halo for 35 convolutions.
                output=np.empty((h*4,w*4,3),dtype=np.uint8)
                for y in range(0,h,a.tile):
                    for x in range(0,w,a.tile):
                        x0,y0=max(0,x-40),max(0,y-40); x1,y1=min(w,x+a.tile+40),min(h,y+a.tile+40)
                        tile=torch.from_numpy(source[y0:y1,x0:x1].copy()).permute(2,0,1).float().unsqueeze(0).to('cuda')/255
                        enhanced=model(tile).clamp_(0,1).squeeze(0).permute(1,2,0).mul_(255).round_().byte().cpu().numpy()
                        ex,ey=min(w,x+a.tile),min(h,y+a.tile)
                        output[y*4:ey*4,x*4:ex*4]=enhanced[(y-y0)*4:(ey-y0)*4,(x-x0)*4:(ex-x0)*4]
                frame=Image.fromarray(output)
                if (ow,oh)!=(w*4,h*4): frame=frame.resize((ow,oh),Image.Resampling.LANCZOS)
                encode.stdin.write(frame.tobytes()); count+=1
                if count%5==0: print(json.dumps({'stage':f'AI 超分 {count}/{len(timestamps)} 帧'}),flush=True)
        encode.stdin.close()
        if decode.wait(timeout=30)!=0 or encode.wait(timeout=120)!=0: raise RuntimeError('Video codec failed')
    finally:
        for process in [decode,encode]:
            if process.poll() is None: process.kill(); process.wait()
    if count!=len(timestamps) or count==0: raise RuntimeError('Output frame count mismatch')
    subprocess.run(['ffmpeg','-v','error','-filter_threads','1','-filter_complex_threads','1','-y','-protocol_whitelist','file,pipe','-i',temp,'-protocol_whitelist','file,pipe','-i',a.input,'-map','0:v:0','-map','1:a?','-map_metadata','1','-metadata:s:v:0','rotate=0','-c:v','copy','-c:a','aac','-b:a','192k','-movflags','+faststart',a.output],check=True,timeout=120)
    actual=probe(a.output); stream=next(s for s in actual['streams'] if s['codec_type']=='video')
    if int(stream['nb_frames'])!=count or abs(float(actual['format']['duration'])-duration)>.2: raise RuntimeError('Output timing mismatch')
    if any(s['codec_type']=='audio' for s in actual['streams']) != any(s['codec_type']=='audio' for s in info['streams']): raise RuntimeError('Output audio stream missing')
    os.unlink(temp)
    print(json.dumps({'result':{'width':ow,'height':oh,'durationSeconds':float(actual['format']['duration']),'fps':float(fps),'frames':count,'audio':any(s['codec_type']=='audio' for s in actual['streams']),'engine':'Real-ESRGAN','model':'realesr-general-x4v3'}}),flush=True)
if __name__=='__main__': main()