"""Single Comfy job: detect -> bounded LTX cleanup -> verify -> literal captions.

No nested HTTP jobs, raw-video output nodes, dynamic package installs, or audio
replacement. Expansion keeps samplers in Comfy's normal cancellation scheduler.
"""
from __future__ import annotations

import json
import math
import os
from pathlib import Path
import threading
import time
import uuid

from . import caption_core as core
from . import literal_renderer as renderer
from . import vision_verifier as vision
from .sttn_repair import H3STTNRepair
from .generation_stage import H3GenerationModelStage

MODEL_NAMES = {
    'unet': 'ltx-2.3-22b-distilled-1.1_transformer_only_fp8_scaled.safetensors',
    'lora': 'h3-rh-local-20261009-8a5f/ltx23_inpaint_rank128_v1_10000steps.safetensors',
    'clip': 'h3-rh-local-20261009-8a5f/gemma_3_12B_it_fp8_e4m3fn.safetensors',
    'projection': 'ltx-2.3_text_projection_bf16.safetensors',
    'video_vae': 'LTX23_video_vae_bf16.safetensors',
    'audio_vae': 'LTX23_audio_vae_bf16.safetensors',
}
_CONDITION_LOCK=threading.Lock()
_CONDITION_KEY=None
_CONDITION_VALUE=None
REPAIR_ENGINE='sttn-v1'


def clone_conditioning(value):
    import torch
    if isinstance(value,torch.Tensor):return value.detach().to('cpu').clone()
    if isinstance(value,list):return [clone_conditioning(v) for v in value]
    if isinstance(value,tuple):return tuple(clone_conditioning(v) for v in value)
    if isinstance(value,dict):return {k:clone_conditioning(v) for k,v in value.items()}
    if value is None or isinstance(value,(str,int,float,bool)):return value
    core.failure('H3_SUBTITLE_DEPENDENCY','固定字幕清理条件含不支持的数据类型')


def conditioning_key():
    import folder_paths
    result=[]
    for key in ['clip','projection']:
        path=Path(folder_paths.get_full_path_or_raise('text_encoders',MODEL_NAMES[key]))
        stat=path.stat();result.append((str(path),stat.st_size,stat.st_mtime_ns,stat.st_ctime_ns))
    return tuple(result)


class H3InpaintConditioningCache:
    @classmethod
    def INPUT_TYPES(cls):return {'required':{}}
    RETURN_TYPES=('CONDITIONING','CONDITIONING')
    FUNCTION='load'
    CATEGORY='video/h3-verified/internal'

    @classmethod
    def IS_CHANGED(cls):return conditioning_key()

    def load(self):
        global _CONDITION_KEY,_CONDITION_VALUE
        import nodes
        key=conditioning_key()
        with _CONDITION_LOCK:
            if _CONDITION_KEY!=key or _CONDITION_VALUE is None:
                # These are fixed inpainting prompts, never user dialogue. Keep
                # detached CPU tensors rather than the 13B encoder/model object.
                clip=nodes.DualCLIPLoader().load_clip(MODEL_NAMES['clip'],MODEL_NAMES['projection'],'ltxv',device='cpu')[0]
                positive=nodes.CLIPTextEncode().encode(clip,'')[0]
                negative=nodes.CLIPTextEncode().encode(clip,'text, subtitles, watermark')[0]
                value=clone_conditioning((positive,negative))
                _CONDITION_VALUE=value;_CONDITION_KEY=key
                del clip,positive,negative
            # Guide nodes add per-video metadata. Never let it mutate the cache
            # or leak conditioning state between two users/tasks.
            return clone_conditioning(_CONDITION_VALUE)


class H3CaptionWarmup:
    @classmethod
    def INPUT_TYPES(cls):return {'required':{}}
    RETURN_TYPES=()
    FUNCTION='warm'
    OUTPUT_NODE=True
    CATEGORY='video/h3-verified/internal'

    @classmethod
    def IS_CHANGED(cls):return float('nan')

    def warm(self):
        import comfy.model_management
        resource_check()
        if REPAIR_ENGINE=='sttn-v1':
            from .sttn_repair import load_model
            load_model()
        else:
            H3InpaintConditioningCache().load()
        # This output-only maintenance job is submitted before opening a new
        # worker to production, never concurrently with a generation on it.
        comfy.model_management.unload_all_models()
        return {'ui':{'text':['H3 fixed inpaint conditioning cache ready']},'result':()}


def interrupt():
    from comfy.model_management import throw_exception_if_processing_interrupted
    throw_exception_if_processing_interrupted()


def audit_logger():
    identifier=uuid.uuid4().hex
    def record(event):
        # Coordinates and enum decisions only; never log image pixels, dialogue,
        # model prompts or credentials. Preserve diagnostics even if a gate fails.
        print('H3_REGION_AUDIT '+json.dumps({'id':identifier,**event}),flush=True)
    return record


def array_frames(images):
    from .frame_pixels import uint8_frames
    return uint8_frames(images, interrupt)


def font_supports(text):
    from fontTools.ttLib import TTFont
    path = renderer.font_path()
    with TTFont(str(path), fontNumber=0, lazy=True) as font:
        cmap = font.getBestCmap()
        if any(not c.isspace() and ord(c) not in cmap for c in text):
            core.failure('H3_SUBTITLE_FONT', '字幕含字体不支持的字符，已停止以避免方框乱码')


def resource_check():
    import folder_paths
    vision.validate_assets()
    if REPAIR_ENGINE=='sttn-v1':
        from .sttn_repair import model_path
        model_path()
        return
    if REPAIR_ENGINE!='ltx-v1':
        core.failure('H3_SUBTITLE_DEPENDENCY','字幕修复引擎配置无效')
    for group,key in [('diffusion_models','unet'),('loras','lora'),('text_encoders','clip'),
                      ('text_encoders','projection'),('vae','video_vae'),('vae','audio_vae')]:
        if not folder_paths.get_full_path(group, MODEL_NAMES[key]):
            core.failure('H3_SUBTITLE_DEPENDENCY', '缺少已验证的字幕清理模型，未返回未清理视频')


def speech_cues(text, audio, duration):
    if not text:
        return []
    if audio is None:
        core.failure('H3_SUBTITLE_ALIGNMENT', '口播任务缺少原音轨')
    import torch
    import torchaudio
    waveform = audio.get('waveform')
    rate = audio.get('sample_rate')
    if not isinstance(waveform,torch.Tensor) or waveform.ndim != 3 or waveform.shape[0] != 1 or not isinstance(rate,int) or rate < 8000:
        core.failure('H3_SUBTITLE_ALIGNMENT', '原音轨格式不支持，未猜测字幕时间')
    if not torch.isfinite(waveform).all() or abs(waveform.shape[-1]/rate-duration) > .4:
        core.failure('H3_SUBTITLE_ALIGNMENT','原音轨时长与视频不一致')
    # ASR receives a new CPU/mono buffer. The original AUDIO input bypasses this
    # node and remains wired directly to the final video writer.
    samples = waveform.detach().to(device='cpu',dtype=torch.float32).mean(dim=1)[0]
    samples = torchaudio.functional.resample(samples,rate,16000).numpy()
    return core.align_literal_cues(text, core.transcribe_words(samples,interrupt),duration)


class H3CaptionPipeline:
    @classmethod
    def INPUT_TYPES(cls):
        return {'required': {'images':('IMAGE',), 'fps':('FLOAT',{'default':24.0,'min':1.0,'max':60.0}),
                'subtitle_text':('STRING',{'multiline':True}), 'seed':('INT',{'default':42,'min':0,'max':2**64-1})},
                'optional': {'audio':('AUDIO',),'classifier_clip':('CLIP',)}}
    RETURN_TYPES=('IMAGE',)
    FUNCTION='process'
    CATEGORY='video/h3-verified'

    @classmethod
    def VALIDATE_INPUTS(cls, subtitle_text):
        try:
            import importlib.util
            resource_check()
            for package in ['rapidocr_onnxruntime','onnxruntime','opencc','fontTools','faster_whisper']:
                if importlib.util.find_spec(package) is None:
                    core.failure('H3_SUBTITLE_DEPENDENCY','字幕处理组件未就绪')
            if subtitle_text:
                font_supports(subtitle_text)
                model=Path(os.environ.get('H3_SUBTITLE_ASR_MODEL','/nonexistent'))
                if not all((model/name).is_file() for name in ['model.bin','config.json','tokenizer.json']):
                    core.failure('H3_SUBTITLE_DEPENDENCY','离线口播模型未就绪')
            return True
        except (ValueError, ImportError, OSError) as exc:
            return str(exc) if str(exc).startswith('H3_') else 'H3_SUBTITLE_DEPENDENCY: 字幕处理依赖未就绪'

    def process(self, images, fps, subtitle_text, seed, audio=None, classifier_clip=None):
        import torch
        from comfy_execution.graph_utils import GraphBuilder
        interrupt()
        if not math.isfinite(fps) or fps <= 0 or not isinstance(subtitle_text,str) or len(subtitle_text)>4096:
            core.failure('H3_SUBTITLE_INPUT','字幕文案或帧率无效')
        original = images.detach().to(device='cpu',dtype=torch.float32)
        pixels = array_frames(original)
        if subtitle_text:
            font_supports(subtitle_text)
        cues = speech_cues(subtitle_text,audio,len(original)/fps)
        check_started=time.monotonic()
        audit=audit_logger()
        classify=lambda samples,interrupt:vision.classify_many(samples,interrupt,clip=classifier_clip,audit=audit)
        review=lambda groups,interrupt:vision.review_many(groups,interrupt,clip=classifier_clip,audit=audit)
        plan = core.caption_plan(pixels,fps,interrupt=interrupt,classify_many=classify,audit=audit,review_many=review)
        crop_json=json.dumps(core.context_crop(plan))
        graph = GraphBuilder()
        accumulated = original
        if plan['regions'] and REPAIR_ENGINE=='sttn-v1':
            resource_check()
            masks=torch.from_numpy(core.mask_from_plan(plan,pixels))
            repaired=graph.node('H3STTNRepair',images=original,masks=masks)
            accumulated=repaired.out(0)
        elif plan['regions']:
            resource_check()
            masks = torch.from_numpy(core.mask_from_plan(plan,pixels))
            model=graph.node('UNETLoader',unet_name=MODEL_NAMES['unet'],weight_dtype='fp8_e4m3fn_fast')
            lora=graph.node('LoraLoaderModelOnly',model=model.out(0),lora_name=MODEL_NAMES['lora'],strength_model=1.0)
            # Native attention: the source workflow's SageAttention failed on
            # this runtime. The successful isolated proof used this exact bypass.
            vae=graph.node('VAELoaderKJ',vae_name=MODEL_NAMES['video_vae'],device='main_device',weight_dtype='bf16')
            audio_vae=graph.node('VAELoaderKJ',vae_name=MODEL_NAMES['audio_vae'],device='main_device',weight_dtype='bf16')
            fixed=graph.node('H3InpaintConditioningCache')
            cond=graph.node('LTXVConditioning',positive=fixed.out(0),negative=fixed.out(1),frame_rate=fps)
            sampler=graph.node('KSamplerSelect',sampler_name='euler_ancestral')
            sigmas=graph.node('ManualSigmas',sigmas='1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0')
            for begin,end,core_begin,core_end in core.chunk_windows(len(original)):
                if not masks[core_begin:core_end].any():
                    continue
                # previous serializes chunks, avoiding parallel intermediate VRAM
                # allocation. Model/CLIP/VAE nodes are shared, not loaded per chunk.
                source=graph.node('H3CaptionChunkSource',images=original,masks=masks,begin=begin,end=end,previous=accumulated,crop_json=crop_json)
                encoded=graph.node('VAEEncode',pixels=source.out(0),vae=vae.out(0))
                guides=graph.node('LTXVAddGuideMulti',positive=cond.out(0),negative=cond.out(1),vae=vae.out(0),latent=encoded.out(0),
                    **{'num_guides':'1','num_guides.image_1':source.out(0),'num_guides.frame_idx_1':0,'num_guides.strength_1':1.0})
                empty_audio=graph.node('LTXVEmptyLatentAudio',frames_number=end-begin,frame_rate=fps,batch_size=1,audio_vae=audio_vae.out(0))
                latent=graph.node('LTXVConcatAVLatent',video_latent=guides.out(2),audio_latent=empty_audio.out(0))
                noise=graph.node('RandomNoise',noise_seed=(seed+core_begin)%(2**64))
                guider=graph.node('CFGGuider',model=lora.out(0),positive=guides.out(0),negative=guides.out(1),cfg=1.0)
                sampled=graph.node('SamplerCustomAdvanced',noise=noise.out(0),guider=guider.out(0),sampler=sampler.out(0),sigmas=sigmas.out(0),latent_image=latent.out(0))
                separated=graph.node('LTXVSeparateAVLatent',av_latent=sampled.out(0))
                cropped=graph.node('LTXVCropGuides',positive=guides.out(0),negative=guides.out(1),latent=separated.out(0))
                decoded=graph.node('VAEDecode',samples=cropped.out(2),vae=vae.out(0))
                mixed=graph.node('H3CaptionBlendChunk',previous=accumulated,original=original,cleaned=decoded.out(0),masks=masks,
                    begin=begin,end=end,core_begin=core_begin,core_end=core_end,crop_json=crop_json)
                accumulated=mixed.out(0)
        final=graph.node('H3CaptionFinalize',images=accumulated,fps=fps,subtitle_text=subtitle_text,
                         cue_json=json.dumps(cues,ensure_ascii=False),plan_json=json.dumps(plan),needs_recheck=bool(plan['regions']))
        summary={'overlays':len(plan['regions']),'semantic_tracks':len(plan.get('semantic_audit',[])),
               'inspection_seconds':time.monotonic()-check_started}
        return {'result':(final.out(0),),'expand':graph.finalize(),'ui':{'h3_verification':[summary]}}


class H3TextRegionVerifier:
    @classmethod
    def INPUT_TYPES(cls):
        return {'required':{'images':('IMAGE',),'regions':('STRING',{'multiline':True,'default':'[]'}),'classifier_clip':('CLIP',)},'optional':{'temporal_review':('BOOLEAN',{'default':False}),'scene_context':('BOOLEAN',{'default':False})}}
    RETURN_TYPES=('STRING',)
    FUNCTION='verify'
    OUTPUT_NODE=True
    CATEGORY='video/h3-verified/internal'

    def verify(self,images,regions,classifier_clip,temporal_review=False,scene_context=False):
        try:rows=json.loads(regions)
        except ValueError:core.failure('H3_SUBTITLE_INPUT','文字区域列表无效')
        if not isinstance(rows,list) or len(rows)>768:core.failure('H3_SUBTITLE_INPUT','文字区域数量无效')
        pixels=array_frames(images);samples=[]
        for row in rows:
            if not isinstance(row,dict) or not isinstance(row.get('frame'),int) or not 0<=row['frame']<len(pixels):
                core.failure('H3_SUBTITLE_INPUT','文字区域帧号无效')
            box=row.get('box')
            if not isinstance(box,list) or len(box)!=4 or any(not isinstance(x,int) for x in box):
                core.failure('H3_SUBTITLE_INPUT','文字区域坐标无效')
            samples.append((pixels[row['frame']],box))
        started=time.monotonic()
        if scene_context and temporal_review:core.failure('H3_SUBTITLE_INPUT','诊断模式不能同时启用')
        if scene_context:labels=vision.classify_scene_many(samples,interrupt,clip=classifier_clip)
        else:labels=vision.review_many([samples],interrupt,clip=classifier_clip) if temporal_review else vision.classify_many(samples,interrupt,clip=classifier_clip)
        result=json.dumps({'labels':labels,'seconds':time.monotonic()-started})
        return {'result':(result,),'ui':{'text':[result]}}


class H3CaptionChunkSource:
    @classmethod
    def INPUT_TYPES(cls):
        return {'required':{'images':('IMAGE',),'masks':('MASK',),'begin':('INT',),'end':('INT',),'previous':('IMAGE',)},'optional':{'crop_json':('STRING',{'default':''})}}
    RETURN_TYPES=('IMAGE',)
    FUNCTION='prepare'
    CATEGORY='video/h3-verified/internal'

    def prepare(self,images,masks,begin,end,previous,crop_json=''):
        import torch
        import torch.nn.functional as F
        interrupt()
        if not 0 <= begin < end <= len(images) or end-begin > 96:
            core.failure('H3_SUBTITLE_INPUT','清理分块边界无效')
        x,y,w,h=json.loads(crop_json) if crop_json else [0,0,images.shape[2],images.shape[1]]
        if not (0<=x<x+w<=images.shape[2] and 0<=y<y+h<=images.shape[1]):
            core.failure('H3_SUBTITLE_INPUT','字幕上下文裁剪越界')
        rgb=images[begin:end,y:y+h,x:x+w].cpu();mask=masks[begin:end,y:y+h,x:x+w].cpu()
        h,w=rgb.shape[1:3];scale=min(1.0,1024/max(h,w));nh=max(32,round(h*scale/32)*32);nw=max(32,round(w*scale/32)*32)
        if (h,w)!=(nh,nw):
            rgb=F.interpolate(rgb.permute(0,3,1,2),size=(nh,nw),mode='bilinear',align_corners=False).permute(0,2,3,1)
            mask=F.interpolate(mask[:,None],size=(nh,nw),mode='nearest')[:,0]
        painted=torch.where(mask[:,:,:,None]>0,torch.ones_like(rgb),rgb)
        pad=(1-len(painted))%4
        if pad:painted=torch.cat([painted[:1].repeat(pad,1,1,1),painted])
        return (painted,)


class H3CaptionBlendChunk:
    @classmethod
    def INPUT_TYPES(cls):
        return {'required':{'previous':('IMAGE',),'original':('IMAGE',),'cleaned':('IMAGE',),'masks':('MASK',),
                'begin':('INT',),'end':('INT',),'core_begin':('INT',),'core_end':('INT',)},'optional':{'crop_json':('STRING',{'default':''})}}
    RETURN_TYPES=('IMAGE',)
    FUNCTION='blend'
    CATEGORY='video/h3-verified/internal'

    def blend(self,previous,original,cleaned,masks,begin,end,core_begin,core_end,crop_json=''):
        import torch
        import torch.nn.functional as F
        interrupt()
        pad=(1-(end-begin))%4
        if len(cleaned) != end-begin+pad or previous.shape != original.shape:
            core.failure('H3_SUBTITLE_FRAME_COUNT','清理前后帧数不一致，未截断原视频')
        cleaned=cleaned[pad:].detach().cpu()
        x,y,w,h=json.loads(crop_json) if crop_json else [0,0,original.shape[2],original.shape[1]]
        if cleaned.shape[1:3] != (h,w):
            cleaned=F.interpolate(cleaned.permute(0,3,1,2),size=(h,w),mode='bilinear',align_corners=False).permute(0,2,3,1)
        mask=masks[core_begin:core_end].cpu();source=original[core_begin:core_end]
        generated=cleaned[core_begin-begin:core_end-begin]
        import numpy as np
        mixed=source.clone()
        for i in range(len(source)):
            interrupt()
            base=(source[i].clamp(0,1).numpy()*255).round().astype(np.uint8)
            replacement=base.copy()
            replacement[y:y+h,x:x+w]=(generated[i].clamp(0,1).numpy()*255).round().astype(np.uint8)
            merged=core.composite_poisson(base,replacement,mask[i].numpy())
            mixed[i]=torch.where(mask[i,:,:,None]>0,torch.from_numpy(merged).to(source.dtype)/255,source[i])
        if not torch.isfinite(mixed).all():core.failure('H3_SUBTITLE_OUTPUT','字幕清理输出含无效像素')
        output=previous.clone();output[core_begin:core_end]=mixed
        return (output,)


class H3CaptionFinalize:
    @classmethod
    def INPUT_TYPES(cls):
        return {'required':{'images':('IMAGE',),'fps':('FLOAT',),'subtitle_text':('STRING',),
                'cue_json':('STRING',),'plan_json':('STRING',),'needs_recheck':('BOOLEAN',)},'optional':{'classifier_clip':('CLIP',)}}
    RETURN_TYPES=('IMAGE',)
    FUNCTION='finish'
    CATEGORY='video/h3-verified'

    def finish(self,images,fps,subtitle_text,cue_json,plan_json,needs_recheck,classifier_clip=None):
        import torch
        import numpy as np
        from PIL import Image
        interrupt();plan=json.loads(plan_json)
        if tuple(images.shape) != (plan['frames'],plan['height'],plan['width'],3):
            core.failure('H3_SUBTITLE_FRAME_COUNT','字幕处理改变了视频尺寸或帧数')
        if needs_recheck:
            audit=audit_logger()
            classify=lambda samples,interrupt:vision.classify_many(samples,interrupt,clip=classifier_clip,audit=audit)
            review=lambda groups,interrupt:vision.review_many(groups,interrupt,clip=classifier_clip,audit=audit)
            after=core.caption_plan(array_frames(images),fps,interrupt=interrupt,classify_many=classify,audit=audit,review_many=review)
            if after['regions']:
                core.failure('H3_SUBTITLE_RESIDUAL','仍检测到原生字幕，未返回双层字幕视频')
        if not subtitle_text:
            return (images,)
        font_supports(subtitle_text)
        cues=renderer.build_cues(subtitle_text,len(images),fps,images.shape[2],images.shape[1],cue_json)
        result=images.clone()
        for i in range(len(images)):
            interrupt()
            active=next((c for c in cues if c['start']<=i/fps<c['end']),None)
            if active is None:continue
            rgb=(images[i].cpu().clamp(0,1).numpy()*255).round().astype(np.uint8)
            drawn=np.array(renderer.render_frame(Image.fromarray(rgb),active['text']))
            changed=torch.from_numpy(np.any(drawn!=rgb,axis=2))[:,:,None]
            # Preserve original floating-point pixels outside the font glyphs.
            result[i]=torch.where(changed,torch.from_numpy(drawn).to(images.dtype)/255,images[i].cpu())
        return (result,)


NODE_CLASS_MAPPINGS={c.__name__:c for c in [H3CaptionPipeline,H3CaptionChunkSource,H3CaptionBlendChunk,H3CaptionFinalize,H3InpaintConditioningCache,H3CaptionWarmup,H3TextRegionVerifier,H3STTNRepair,H3GenerationModelStage]}
NODE_DISPLAY_NAME_MAPPINGS={'H3CaptionPipeline':'H3 统一字幕处理（检测/清理/校验/烧录）'}

from .worker_state import register_worker_state
register_worker_state()