import os
import glob
from typing import Any, Dict, List
import gradio as gr
import spaces

from core.pipelines.workflow_executor import WorkflowExecutor
from core.model_manager import model_manager

def _extract_required_models(workflow: Dict[str, Any]) -> List[str]:
    required_models = []
    MODEL_KEYS = {'unet_name', 'clip_name', 'vae_name', 'model_name', 'ckpt_name', 'clip_name1', 'clip_name2'}
    MODEL_EXTENSIONS = ('.safetensors', '.ckpt', '.pt', '.bin', '.pth')
    
    for node_id, node_data in workflow.items():
        if node_data.get('class_type') == 'LoraLoader':
            continue
        inputs = node_data.get('inputs', {})
        for k, v in inputs.items():
            if isinstance(v, str) and v and k != 'lora_name':
                if k in MODEL_KEYS or any(v.lower().endswith(ext) for ext in MODEL_EXTENSIONS):
                    if v not in required_models and not v.startswith('[') and not v.endswith(']'):
                        required_models.append(v)
    return required_models

def _extract_video_paths(out: Any) -> List[str]:
    video_files = []
    
    def process_item(item):
        if item is None:
            return
        if isinstance(item, (list, tuple)):
            for sub in item:
                process_item(sub)
            return
        if isinstance(item, str) and os.path.exists(item):
            video_files.append(item)
            return
        for attr in ['path', 'saved_path', 'filepath', 'filename', 'full_path']:
            val = getattr(item, attr, None)
            if isinstance(val, str) and os.path.exists(val):
                video_files.append(val)
                return
            elif isinstance(val, str):
                try:
                    import folder_paths
                    possible = os.path.join(folder_paths.get_output_directory(), val)
                    if os.path.exists(possible):
                        video_files.append(possible)
                        return
                except Exception:
                    pass
        if hasattr(item, '__dict__'):
            for k, v in item.__dict__.items():
                if isinstance(v, str) and (v.endswith('.mp4') or v.endswith('.webm') or v.endswith('.gif')):
                    if os.path.exists(v):
                        video_files.append(v)
                        return
                    try:
                        import folder_paths
                        possible = os.path.join(folder_paths.get_output_directory(), v)
                        if os.path.exists(possible):
                            video_files.append(possible)
                            return
                    except Exception:
                        pass
        if isinstance(item, dict):
            ui_info = item.get("ui", {})
            for key in ["images", "videos", "video"]:
                for sub in ui_info.get(key, []):
                    if isinstance(sub, dict) and "filename" in sub:
                        fn = sub["filename"]
                        subfolder = sub.get("subfolder", "")
                        try:
                            import folder_paths
                            full = os.path.join(folder_paths.get_output_directory(), subfolder, fn) if subfolder else os.path.join(folder_paths.get_output_directory(), fn)
                            if os.path.exists(full):
                                video_files.append(full)
                                return
                        except Exception:
                            pass

    process_item(out)

    if not video_files:
        try:
            import folder_paths
            out_dir = folder_paths.get_output_directory()
            mp4_files = glob.glob(os.path.join(out_dir, "**", "*.mp4"), recursive=True)
            if mp4_files:
                mp4_files.sort(key=os.path.getmtime, reverse=True)
                video_files.append(mp4_files[0])
        except Exception:
            pass
            
    return video_files

def _execute_h3_workflow_gpu(workflow: Dict[str, Any]):
    initial_objects = {}
    return WorkflowExecutor.execute_workflow(workflow, initial_objects=initial_objects)

def generate_h3_video_wrapper(process_inputs_func, ui_inputs: dict, progress=gr.Progress(track_tqdm=True)):
    progress(0.1, desc="Assembling H3 Video Workflow...")
    try:
        workflow, extra_data = process_inputs_func(ui_inputs, progress=progress)
    except TypeError:
        workflow, extra_data = process_inputs_func(ui_inputs)
    
    required_models = _extract_required_models(workflow)
    if required_models:
        progress(0.2, desc=f"Ensuring models are downloaded: {len(required_models)} file(s)...")
        model_manager.ensure_models_downloaded(required_models, progress=progress)

    zero_gpu_duration = ui_inputs.get('zero_gpu_duration', 120)
    try:
        duration = int(zero_gpu_duration)
    except (ValueError, TypeError):
        duration = 120
    if duration <= 0:
        duration = 120

    progress(0.4, desc=f"Executing H3 Video Workflow on GPU (ZeroGPU {duration}s)...")
    gpu_runner = spaces.GPU(duration=duration)(_execute_h3_workflow_gpu)
    raw_output = gpu_runner(workflow)
    
    video_files = _extract_video_paths(raw_output)
    final_video = video_files[0] if video_files else None
    print("✅ H3 Video generated successfully")
    
    progress(1.0, desc="H3 Video Generation Complete!")
    return final_video