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):
            if item.lower().endswith(('.mp4', '.webm', '.gif', '.mov', '.mkv')):
                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_ltx_workflow_gpu(workflow: Dict[str, Any]):
    initial_objects = {}
    return WorkflowExecutor.execute_workflow(workflow, initial_objects=initial_objects)

def generate_ltx_video_wrapper(process_inputs_func, ui_inputs: dict, progress=gr.Progress(track_tqdm=True)):
    progress(0.1, desc="Assembling LTX-2.5 Video Workflow...")
    batch_count = int(ui_inputs.get('batch_count', 1))
    all_video_files = []

    for b_idx in range(batch_count):
        batch_msg = f" (Batch {b_idx + 1}/{batch_count})" if batch_count > 1 else ""
        current_ui_inputs = ui_inputs.copy()
        
        orig_seed = int(current_ui_inputs.get('seed', -1))
        if orig_seed != -1 and b_idx > 0:
            current_ui_inputs['seed'] = orig_seed + b_idx

        try:
            workflow, extra_data = process_inputs_func(current_ui_inputs, progress=progress)
        except TypeError:
            workflow, extra_data = process_inputs_func(current_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 = current_ui_inputs.get('zero_gpu_duration', 60)
        try:
            duration = int(zero_gpu_duration)
        except (ValueError, TypeError):
            duration = 60
        if duration <= 0:
            duration = 60

        progress(0.4, desc=f"Executing LTX-2.5 Workflow on GPU{batch_msg} (ZeroGPU {duration}s)...")
        gpu_runner = spaces.GPU(duration=duration)(_execute_ltx_workflow_gpu)
        raw_output = gpu_runner(workflow)
        
        batch_videos = _extract_video_paths(raw_output)
        if batch_videos:
            all_video_files.extend(batch_videos)
            print(f"✅ LTX-2.5 Video generated{batch_msg}")
    
    progress(1.0, desc="LTX-2.5 Video Generation Complete!")
    if not all_video_files:
        return None
    return all_video_files[0] if batch_count == 1 else all_video_files

# Alias for backwards compatibility
generate_h3_video_wrapper = generate_ltx_video_wrapper