import os
import random
import requests
import hashlib
import re
from typing import Sequence, Mapping, Any, Union, Set
from pathlib import Path
import shutil

import gradio as gr
from huggingface_hub import hf_hub_download, constants as hf_constants
import torch
import numpy as np
from PIL import Image, ImageChops
import yaml

from core.settings import *

MODELS_ROOT_DIR = "ComfyUI/models"


class UniqueKeyLoader(yaml.SafeLoader):
    """
    A custom YAML loader that handles duplicate keys by grouping their values into a list.
    """
    def construct_mapping(self, node, deep=False):
        mapping = []
        for key_node, value_node in node.value:
            key = self.construct_object(key_node, deep=deep)
            value = self.construct_object(value_node, deep=deep)
            mapping.append((key, value))
        
        result = {}
        for k, v in mapping:
            if k in result:
                if isinstance(result[k], list):
                    result[k].append(v)
                else:
                    result[k] = [result[k], v]
            else:
                result[k] = v
        return result

UniqueKeyLoader.add_constructor(yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, UniqueKeyLoader.construct_mapping)

def save_uploaded_file_with_hash(file_obj: gr.File, target_dir: str) -> str:
    if not file_obj:
        return ""
    
    temp_path = file_obj.name
    
    sha256 = hashlib.sha256()
    with open(temp_path, 'rb') as f:
        for block in iter(lambda: f.read(65536), b''):
            sha256.update(block)
    
    file_hash = sha256.hexdigest()
    _, extension = os.path.splitext(temp_path)
    hashed_filename = f"{file_hash}{extension.lower()}"
    
    dest_path = os.path.join(target_dir, hashed_filename)
    
    os.makedirs(target_dir, exist_ok=True)
    if not os.path.exists(dest_path):
        shutil.copy(temp_path, dest_path)
        print(f"✅ Saved uploaded file as: {dest_path}")
    else:
        print(f"ℹ️ File already exists (deduplicated): {dest_path}")
        
    return hashed_filename

def bytes_to_gb(byte_size: int) -> float:
    if byte_size is None or byte_size == 0:
        return 0.0
    return round(byte_size / (1024 ** 3), 2)

def get_directory_size(path: str) -> int:
    total_size = 0
    if not os.path.exists(path):
        return 0
    try:
        for dirpath, _, filenames in os.walk(path):
            for f in filenames:
                fp = os.path.join(dirpath, f)
                if os.path.isfile(fp) and not os.path.islink(fp):
                    total_size += os.path.getsize(fp)
    except OSError as e:
        print(f"Warning: Could not access {path} to calculate size: {e}")
    return total_size

def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
    try:
        return obj[index]
    except (KeyError, IndexError):
        try:
            return obj["result"][index]
        except (KeyError, IndexError):
            return None

def sanitize_prompt(prompt: str) -> str:
    if not isinstance(prompt, str):
        return ""
    return "".join(char for char in prompt if char.isprintable() or char in ('\n', '\t'))

def sanitize_id(input_id: str) -> str:
    if not isinstance(input_id, str):
        return ""
    input_id = input_id.strip()
    if "civitai" in input_id.lower():
        version_match = re.search(r'modelVersionId=(\d+)', input_id)
        if version_match:
            return version_match.group(1)
        model_match = re.search(r'/models/(\d+)', input_id)
        if model_match:
            return model_match.group(1)
    return re.sub(r'[^0-9]', '', input_id)

def sanitize_url(url: str) -> str:
    if not isinstance(url, str):
        raise ValueError("URL must be a string.")
    url = url.strip()
    if not re.match(r'^https?://[^\s/$.?#].[^\s]*$', url):
        raise ValueError("Invalid URL format or scheme. Only HTTP and HTTPS are allowed.")
    return url

def sanitize_filename(filename: str) -> str:
    if not isinstance(filename, str):
        return ""
    sanitized = filename.replace('..', '')
    sanitized = re.sub(r'[^\w\.\-]', '_', sanitized)
    return sanitized.lstrip('/\\')

def get_civitai_file_info(version_id: str) -> dict | None:
    api_url = f"https://civitai.com/api/v1/model-versions/{version_id}"
    try:
        response = requests.get(api_url, timeout=10)
        response.raise_for_status()
        data = response.json()
        
        model_type = data.get('model', {}).get('type')
        
        result_file = None
        for file_data in data.get('files', []):
            if file_data.get('type') == 'Model' and file_data['name'].endswith(('.safetensors', '.pt', '.bin')):
                result_file = file_data.copy()
                break
        
        if not result_file and data.get('files'):
            result_file = data['files'][0].copy()
            
        if result_file:
            result_file['model_type'] = model_type
            return result_file
    except Exception:
        return None

def download_file(url: str, save_path: str, api_key: str = None, progress=None, desc: str = "") -> str:
    if os.path.exists(save_path):
        return f"File already exists: {os.path.basename(save_path)}"
    
    headers = {'Authorization': f'Bearer {api_key}'} if api_key and api_key.strip() else {}
    try:
        if progress:
            progress(0, desc=desc)
        
        response = requests.get(url, stream=True, headers=headers, timeout=15)
        response.raise_for_status()
        total_size = int(response.headers.get('content-length', 0))
        
        with open(save_path, "wb") as f:
            downloaded = 0
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
                if progress and total_size > 0:
                    downloaded += len(chunk)
                    progress(downloaded / total_size, desc=desc)
        return f"Successfully downloaded: {os.path.basename(save_path)}"
    except Exception as e:
        if os.path.exists(save_path):
            os.remove(save_path)
        return f"Download failed for {os.path.basename(save_path)}: {e}"

def get_lora_path(source: str, id_or_url: str, civitai_key: str, progress) -> tuple[str | None, str]:
    if not id_or_url or not id_or_url.strip():
        return None, "No ID/URL provided."

    try:
        if source == "Civitai":
            version_id = sanitize_id(id_or_url)
            if not version_id:
                return None, "Invalid Civitai ID provided. Must be numeric."
            
            file_info = get_civitai_file_info(version_id)
            if file_info:
                model_type = file_info.get('model_type')
                if model_type and model_type.lower() == 'checkpoint':
                    return None, f"Invalid Civitai model type '{model_type}' for LoRA. Checkpoint models are not allowed."
            
            filename = sanitize_filename(f"civitai_{version_id}.safetensors")
            local_path = os.path.join(LORA_DIR, filename)
            api_key_to_use = civitai_key
            source_name = f"Civitai ID {version_id}"
        elif source == "Hugging Face":
            parts = id_or_url.strip().split('/')
            if len(parts) < 3:
                return None, "Invalid Hugging Face path. Format: repo_owner/repo_name/filename"
            repo_id = f"{parts[0]}/{parts[1]}"
            repo_file_path = "/".join(parts[2:])
            unique_name = id_or_url.strip().replace('/', '_')
            filename = sanitize_filename(unique_name)
            local_path = os.path.join(LORA_DIR, filename)
            source_name = f"HF {repo_file_path}"
        else:
            return None, "Invalid source."

    except ValueError as e:
        return None, f"Input validation failed: {e}"

    if os.path.lexists(local_path):
        if not os.path.exists(local_path):
            os.remove(local_path)
        else:
            return local_path, "File already exists."

    if source == "Civitai":
        if not file_info or not file_info.get('downloadUrl'):
            return None, f"Could not get download link for {source_name}."

        status = download_file(file_info['downloadUrl'], local_path, api_key_to_use, progress=progress, desc=f"Downloading {source_name}")
        return (local_path, status) if "Successfully" in status else (None, status)
    elif source == "Hugging Face":
        try:
            if progress and callable(progress): progress(0, desc=f"Downloading {source_name}")
            cached_path = hf_hub_download(repo_id=repo_id, filename=repo_file_path, token=os.environ.get("HF_TOKEN"))
            os.makedirs(LORA_DIR, exist_ok=True)
            if os.path.lexists(local_path):
                if not os.path.exists(local_path):
                    try:
                        os.remove(local_path)
                    except OSError:
                        pass
            if not os.path.exists(local_path):
                try:
                    os.symlink(cached_path, local_path)
                except (OSError, NotImplementedError):
                    shutil.copyfile(cached_path, local_path)
            if progress and callable(progress): progress(1.0, desc=f"Downloaded {source_name}")
            return local_path, f"Successfully downloaded: {filename}"
        except Exception as e:
            return None, f"Hugging Face download failed: {e}"


def _ensure_model_downloaded(display_name: str, progress=gr.Progress()):
    if display_name not in ALL_MODEL_MAP:
        for cat_dir in CATEGORY_TO_DIR_MAP.values():
            check_path = os.path.join(cat_dir, display_name)
            if os.path.exists(check_path):
                return display_name
        raise ValueError(f"Model '{display_name}' not found in configuration.")

    model_info = ALL_MODEL_MAP[display_name]
    repo_filename = model_info[1]
    base_filename = os.path.basename(repo_filename)

    download_info = ALL_FILE_DOWNLOAD_MAP.get(base_filename)
    if not download_info:
        raise gr.Error(f"Model '{base_filename}' not found in file_list.yaml. Cannot download.")

    category = download_info.get("category")
    dest_dir = CATEGORY_TO_DIR_MAP.get(category)
    
    if not dest_dir:
        raise ValueError(f"Unknown YAML category '{category}' for '{base_filename}'.")
    
    dest_path = os.path.join(dest_dir, base_filename)

    if os.path.lexists(dest_path):
        if not os.path.exists(dest_path):
            print(f"⚠️ Found and removed broken symlink: {dest_path}")
            os.remove(dest_path)
        else:
            return base_filename

    source = download_info.get("source")
    try:
        progress(0, desc=f"Downloading: {base_filename}")
        
        if source == "hf":
            repo_id = download_info.get("repo_id")
            hf_filename = download_info.get("repository_file_path", base_filename)
            if not repo_id:
                raise ValueError(f"repo_id is missing for HF model '{base_filename}'")
            
            cached_path = hf_hub_download(repo_id=repo_id, filename=hf_filename, token=os.environ.get("HF_TOKEN"))
            os.makedirs(dest_dir, exist_ok=True)
            os.symlink(cached_path, dest_path)
            print(f"✅ Symlinked '{cached_path}' to '{dest_path}'")

        elif source == "civitai":
            model_version_id = download_info.get("model_version_id")
            if not model_version_id:
                raise ValueError(f"model_version_id is missing for Civitai model '{base_filename}'")
            
            file_info = get_civitai_file_info(model_version_id)
            if not file_info or not file_info.get('downloadUrl'):
                raise ConnectionError(f"Could not get download URL for Civitai model version ID {model_version_id}")
            
            status = download_file(
                file_info['downloadUrl'], dest_path, api_key=os.environ.get("CIVITAI_API_KEY", ""), progress=progress, desc=f"Downloading: {base_filename}"
            )
            if "Failed" in status:
                raise ConnectionError(status)
        else:
            raise NotImplementedError(f"Download source '{source}' is not implemented for '{base_filename}'")
            
        progress(1.0, desc=f"Downloaded: {base_filename}")

    except Exception as e:
        if os.path.lexists(dest_path):
            try:
                os.remove(dest_path)
            except OSError: pass
        raise gr.Error(f"Failed to download and link '{display_name}': {e}")
    
    return base_filename


def ensure_file_downloaded(filename: str, progress=None):
    if not filename or filename == "None":
        return

    download_info = ALL_FILE_DOWNLOAD_MAP.get(filename)
    if not download_info:
        print(f"⚠️ Warning: File '{filename}' not found in configuration (file_list.yaml). Cannot download.")
        return

    category = download_info.get("category", "loras")
    dest_dir = CATEGORY_TO_DIR_MAP.get(category, LORA_DIR)
    dest_path = os.path.join(dest_dir, filename)

    if os.path.lexists(dest_path):
        if not os.path.exists(dest_path):
            print(f"⚠️ Found and removed broken symlink: {dest_path}")
            os.remove(dest_path)
        else:
            return

    source = download_info.get("source")
    try:
        if source == "hf":
            repo_id = download_info.get("repo_id")
            repo_filename = download_info.get("repository_file_path", filename)
            if not repo_id:
                raise ValueError("repo_id is missing for Hugging Face download.")

            if progress and callable(progress):
                progress(0, desc=f"Downloading: {filename}")
            cached_path = hf_hub_download(repo_id=repo_id, filename=repo_filename, token=os.environ.get("HF_TOKEN"))
            os.makedirs(dest_dir, exist_ok=True)
            os.symlink(cached_path, dest_path)
            print(f"✅ Symlinked '{cached_path}' to '{dest_path}'")
            if progress and callable(progress):
                progress(1.0, desc=f"Downloaded: {filename}")

        elif source == "civitai":
            model_version_id = download_info.get("model_version_id")
            if not model_version_id:
                raise ValueError("model_version_id is missing for Civitai download.")

            file_info = get_civitai_file_info(model_version_id)
            if not file_info or not file_info.get('downloadUrl'):
                raise ConnectionError(f"Could not get download URL for Civitai model version ID {model_version_id}")

            status = download_file(
                file_info['downloadUrl'],
                dest_path,
                api_key=os.environ.get("CIVITAI_API_KEY", ""),
                progress=progress,
                desc=f"Downloading: {filename}"
            )
            if "Failed" in status:
                raise ConnectionError(status)
        else:
            raise NotImplementedError(f"Download source '{source}' is not implemented for '{filename}'.")

    except Exception as e:
        if os.path.lexists(dest_path):
            try:
                os.remove(dest_path)
            except OSError:
                pass
        raise gr.Error(f"Failed to download file '{filename}': {e}")


def get_model_generation_defaults(model_display_name: str, model_type: str, defaults_config: dict):
    final_defaults = {
        'steps': 25, 'cfg': 7.0, 'sampler_name': 'euler', 'scheduler': 'simple',
        'positive_prompt': '', 'negative_prompt': ''
    }

    if 'Default' in defaults_config:
        final_defaults.update(defaults_config['Default'])

    model_type_key = next((key for key in defaults_config if key.lower().replace(" ", "-").replace(".", "") == model_type.lower()), None)
    if model_type_key:
        model_type_config = defaults_config[model_type_key]
        if '_defaults' in model_type_config:
            final_defaults.update(model_type_config['_defaults'])
            
        if model_display_name in model_type_config:
            final_defaults.update(model_type_config[model_display_name])

    return final_defaults

def get_filename_prefix() -> str:
    import time
    return f"H3_{int(time.time())}"

def save_temp_image(img):
    if img is None:
        return None
    _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    input_dir = os.path.join(_PROJECT_ROOT, "input")
    os.makedirs(input_dir, exist_ok=True)
    if isinstance(img, Image.Image):
        filename = f"temp_image_{random.randint(10000, 99999)}.png"
        filepath = os.path.join(input_dir, filename)
        img.save(filepath, "PNG")
        return os.path.basename(filepath)
    elif isinstance(img, str):
        if not img:
            return None
        if os.path.exists(img):
            ext = os.path.splitext(img)[1] or ".png"
            filename = f"temp_image_{random.randint(10000, 99999)}{ext}"
            save_path = os.path.join(input_dir, filename)
            shutil.copy(img, save_path)
            return os.path.basename(save_path)
        if os.path.exists(os.path.join(input_dir, img)):
            return img
        return os.path.basename(img)
    return None

def save_temp_audio(audio_path):
    if not audio_path:
        return None
    _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    input_dir = os.path.join(_PROJECT_ROOT, "input")
    os.makedirs(input_dir, exist_ok=True)
    if os.path.exists(audio_path):
        ext = os.path.splitext(audio_path)[1] or ".wav"
        filename = f"temp_audio_{random.randint(10000, 99999)}{ext}"
        save_path = os.path.join(input_dir, filename)
        shutil.copy(audio_path, save_path)
        return os.path.basename(filename)
    if os.path.exists(os.path.join(input_dir, audio_path)):
        return audio_path
    return os.path.basename(audio_path)

def save_temp_video(video_path):
    if not video_path:
        return None
    _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    input_dir = os.path.join(_PROJECT_ROOT, "input")
    os.makedirs(input_dir, exist_ok=True)
    if os.path.exists(video_path):
        ext = os.path.splitext(video_path)[1] or ".mp4"
        filename = f"temp_video_{random.randint(10000, 99999)}{ext}"
        save_path = os.path.join(input_dir, filename)
        shutil.copy(video_path, save_path)
        return os.path.basename(filename)
    if os.path.exists(os.path.join(input_dir, video_path)):
        return video_path
    return os.path.basename(video_path)

def handle_seed(seed_value: int, max_val: int = 2**32 - 1) -> int:
    if seed_value == -1 or seed_value is None:
        return random.randint(0, max_val)
    return int(seed_value)

def process_lora_inputs(ui_values: dict, prefix: str = "", progress=None) -> list:
    active_loras_for_gpu = []

    # 1. Check direct prefix format (e.g. lora_sources_h3_fl2va)
    lora_sources = ui_values.get(f'lora_sources_{prefix}', []) if prefix else []
    lora_ids = ui_values.get(f'lora_ids_{prefix}', []) if prefix else []
    lora_scales = ui_values.get(f'lora_scales_{prefix}', []) if prefix else []

    if isinstance(lora_sources, list) and isinstance(lora_ids, list):
        for source, val, scale in zip(lora_sources, lora_ids, lora_scales):
            scale_val = float(scale) if scale is not None else 1.0
            if scale_val > 0 and val and str(val).strip():
                lora_id = str(val).strip()
                lora_filename = None
                if source == "File":
                    lora_filename = sanitize_filename(lora_id)
                    local_path = os.path.join(LORA_DIR, lora_filename)
                    if not os.path.exists(local_path):
                        raise gr.Error(f"Uploaded LoRA file '{lora_id}' no longer exists on server. Please re-upload it.")
                elif source in ("Civitai", "Hugging Face"):
                    local_path, status = get_lora_path(source, lora_id, os.environ.get("CIVITAI_API_KEY", ""), progress)
                    if local_path:
                        lora_filename = os.path.basename(local_path)
                    else:
                        raise gr.Error(f"Failed to prepare LoRA {lora_id}: {status}")

                if lora_filename:
                    active_loras_for_gpu.append({
                        "lora_name": lora_filename,
                        "strength_model": scale_val,
                        "strength_clip": scale_val
                    })

    # 2. Check lora_data flat list format (e.g. [source1, id1, scale1, upload1, ...])
    lora_data = ui_values.get('lora_data', [])
    if lora_data and not active_loras_for_gpu:
        sources, ids, scales, files = lora_data[0::4], lora_data[1::4], lora_data[2::4], lora_data[3::4]
        for source, lora_id, scale, _ in zip(sources, ids, scales, files):
            scale_val = float(scale) if scale is not None else 1.0
            if scale_val > 0 and lora_id and str(lora_id).strip():
                lora_id_str = str(lora_id).strip()
                lora_filename = None
                if source == "File":
                    lora_filename = sanitize_filename(lora_id_str)
                    local_path = os.path.join(LORA_DIR, lora_filename)
                    if not os.path.exists(local_path):
                        raise gr.Error(f"Uploaded LoRA file '{lora_id_str}' no longer exists on server. Please re-upload it.")
                elif source in ("Civitai", "Hugging Face"):
                    local_path, status = get_lora_path(source, lora_id_str, os.environ.get("CIVITAI_API_KEY", ""), progress)
                    if local_path:
                        lora_filename = os.path.basename(local_path)
                    else:
                        raise gr.Error(f"Failed to prepare LoRA {lora_id_str}: {status}")

                if lora_filename:
                    active_loras_for_gpu.append({
                        "lora_name": lora_filename,
                        "strength_model": scale_val,
                        "strength_clip": scale_val
                    })

    # 3. Check direct 'loras' list of dicts (from MCP or custom payload)
    raw_loras = ui_values.get('loras', [])
    if raw_loras and not active_loras_for_gpu and isinstance(raw_loras, list):
        for item in raw_loras:
            if isinstance(item, dict):
                if "lora_name" in item:
                    active_loras_for_gpu.append(item)
                else:
                    src = item.get("source", "Hugging Face")
                    val = item.get("lora_value") or item.get("id_or_url") or item.get("lora_id")
                    scale = item.get("scale", 1.0)
                    scale_val = float(scale) if scale is not None else 1.0
                    if scale_val > 0 and val and str(val).strip():
                        lora_id_str = str(val).strip()
                        lora_filename = None
                        if src == "File":
                            lora_filename = sanitize_filename(lora_id_str)
                            local_path = os.path.join(LORA_DIR, lora_filename)
                            if not os.path.exists(local_path):
                                raise gr.Error(f"Uploaded LoRA file '{lora_id_str}' no longer exists on server. Please re-upload it.")
                        elif src in ("Civitai", "Hugging Face"):
                            local_path, status = get_lora_path(src, lora_id_str, os.environ.get("CIVITAI_API_KEY", ""), progress)
                            if local_path:
                                lora_filename = os.path.basename(local_path)
                            else:
                                raise gr.Error(f"Failed to prepare LoRA {lora_id_str}: {status}")
                        if lora_filename:
                            active_loras_for_gpu.append({
                                "lora_name": lora_filename,
                                "strength_model": scale_val,
                                "strength_clip": scale_val
                            })

    return active_loras_for_gpu


def load_h3_controlnet_config():
    _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    _CN_MODEL_LIST_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'h3_controlnet_models.yaml')
    try:
        with open(_CN_MODEL_LIST_PATH, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        return config.get("H3_ControlNet", {}) if isinstance(config, dict) else {}
    except Exception as e:
        print(f"Error loading h3_controlnet_models.yaml: {e}")
        return {}


def get_h3_cn_defaults(arch_val="MiniMax-H3"):
    cn_full_config = load_h3_controlnet_config()
    cn_config = cn_full_config.get(arch_val, [])
    if not cn_config and cn_full_config:
        cn_config = next(iter(cn_full_config.values()), [])

    if not cn_config:
        return ["Canny", "Depth", "HED", "MLSD", "Pose"], "Canny", ["alibaba-pai/MiniMax-H3-Fun-Controlnet-Union"], "alibaba-pai/MiniMax-H3-Fun-Controlnet-Union", "minimax_h3_fun_controlnet_union_pruned_int8_convrot.safetensors"

    all_types = []
    for model in cn_config:
        for t in model.get("Type", []):
            if t not in all_types:
                all_types.append(t)
    default_type = all_types[0] if all_types else "Canny"

    series_choices = []
    if default_type:
        for model in cn_config:
            if default_type in model.get("Type", []):
                s = model.get("Series", "Default")
                if s not in series_choices:
                    series_choices.append(s)
    default_series = series_choices[0] if series_choices else ""

    filepath = ""
    if default_series and default_type:
        for model in cn_config:
            if model.get("Series") == default_series and default_type in model.get("Type", []):
                filepath = model.get("Filepath", "")
                break

    return all_types, default_type, series_choices, default_series, filepath


def process_h3_controlnet_inputs(ui_values: dict, prefix: str = "", progress=None) -> list:
    active_cns = []

    # 1. Check direct prefix format (e.g. h3_controlnet_videos_h3_fl2va)
    videos = ui_values.get(f'h3_controlnet_videos_{prefix}', []) if prefix else []
    types = ui_values.get(f'h3_controlnet_types_{prefix}', []) if prefix else []
    series = ui_values.get(f'h3_controlnet_series_{prefix}', []) if prefix else []
    strengths = ui_values.get(f'h3_controlnet_strengths_{prefix}', []) if prefix else []
    start_percents = ui_values.get(f'h3_controlnet_start_percents_{prefix}', []) if prefix else []
    end_percents = ui_values.get(f'h3_controlnet_end_percents_{prefix}', []) if prefix else []
    filepaths = ui_values.get(f'h3_controlnet_filepaths_{prefix}', []) if prefix else []

    if isinstance(videos, list) and isinstance(types, list):
        for idx, vid in enumerate(videos):
            if vid:
                saved_vid_name = save_temp_video(vid) if isinstance(vid, str) and os.path.exists(vid) else vid
                if not saved_vid_name:
                    continue

                fp = filepaths[idx] if idx < len(filepaths) and filepaths[idx] and filepaths[idx] != "None" else None
                if not fp:
                    fp = "minimax_h3_fun_controlnet_union_pruned_int8_convrot.safetensors"

                ensure_file_downloaded(fp, progress=progress)

                st = float(strengths[idx]) if idx < len(strengths) and strengths[idx] is not None else 1.0
                sp = float(start_percents[idx]) if idx < len(start_percents) and start_percents[idx] is not None else 0.0
                ep = float(end_percents[idx]) if idx < len(end_percents) and end_percents[idx] is not None else 1.0

                active_cns.append({
                    "video": saved_vid_name,
                    "control_net_name": fp,
                    "strength": st,
                    "start_percent": sp,
                    "end_percent": ep
                })

    # 2. Check flat component list (e.g. [video1, type1, series1, strength1, start1, end1, filepath1, ...])
    cn_data = ui_values.get(f'h3_controlnet_data_{prefix}', []) or ui_values.get('h3_controlnet_data', [])
    if cn_data and not active_cns and isinstance(cn_data, list):
        stride = 7
        for i in range(0, len(cn_data), stride):
            chunk = cn_data[i:i+stride]
            if len(chunk) >= 1 and chunk[0]:
                vid = chunk[0]
                saved_vid_name = save_temp_video(vid) if isinstance(vid, str) and os.path.exists(vid) else vid
                if not saved_vid_name:
                    continue
                fp = chunk[6] if len(chunk) > 6 and chunk[6] and chunk[6] != "None" else "minimax_h3_fun_controlnet_union_pruned_int8_convrot.safetensors"
                ensure_file_downloaded(fp, progress=progress)
                st = float(chunk[3]) if len(chunk) > 3 and chunk[3] is not None else 1.0
                sp = float(chunk[4]) if len(chunk) > 4 and chunk[4] is not None else 0.0
                ep = float(chunk[5]) if len(chunk) > 5 and chunk[5] is not None else 1.0
                active_cns.append({
                    "video": saved_vid_name,
                    "control_net_name": fp,
                    "strength": st,
                    "start_percent": sp,
                    "end_percent": ep
                })

    # 3. Check direct 'h3_controlnets' list of dicts (from MCP or custom payload)
    raw_cns = ui_values.get('h3_controlnets', [])
    if raw_cns and not active_cns and isinstance(raw_cns, list):
        for item in raw_cns:
            if isinstance(item, dict):
                vid = item.get('video') or item.get('control_video') or item.get('file')
                if vid:
                    saved_vid_name = save_temp_video(vid) if isinstance(vid, str) and os.path.exists(vid) else vid
                    if not saved_vid_name:
                        continue
                    fp = item.get('control_net_name') or item.get('filepath') or item.get('name') or "minimax_h3_fun_controlnet_union_pruned_int8_convrot.safetensors"
                    ensure_file_downloaded(fp, progress=progress)
                    active_cns.append({
                        "video": saved_vid_name,
                        "control_net_name": fp,
                        "strength": float(item.get('strength', 1.0)),
                        "start_percent": float(item.get('start_percent', 0.0)),
                        "end_percent": float(item.get('end_percent', 1.0))
                    })

    return active_cns


def process_h3_guide_inputs(ui_values: dict, prefix: str = "", progress=None) -> list:
    """
    Parses and prepares MiniMax H3 keyframe guide inputs.
    Supports:
      1. Direct prefix format (e.g. h3_guide_images_h3_ref2va, h3_guide_times_h3_ref2va, etc.)
      2. Direct 'h3_guides' or 'guides' list of dictionaries (from MCP or custom payload)
    Returns:
      List of guide dicts: [{'frame_idx': int, 'image': str, 'video': str, 'audio': str}, ...]
    """
    active_guides = []

    # 1. Direct prefix format from Gradio UI
    images = ui_values.get(f'h3_guide_images_{prefix}', []) if prefix else []
    videos = ui_values.get(f'h3_guide_videos_{prefix}', []) if prefix else []
    audios = ui_values.get(f'h3_guide_audios_{prefix}', []) if prefix else []
    times = ui_values.get(f'h3_guide_times_{prefix}', []) if prefix else []
    frames = ui_values.get(f'h3_guide_frames_{prefix}', []) if prefix else []

    if images or videos or audios:
        max_len = max(len(images), len(videos), len(audios), len(times), len(frames))
        for i in range(max_len):
            img = images[i] if i < len(images) else None
            vid = videos[i] if i < len(videos) else None
            aud = audios[i] if i < len(audios) else None
            t_val = times[i] if i < len(times) else None
            f_val = frames[i] if i < len(frames) else None

            saved_img = save_temp_image(img) if img is not None else None
            saved_vid = save_temp_video(vid) if vid else None
            saved_aud = save_temp_audio(aud) if aud else None

            if not (saved_img or saved_vid or saved_aud):
                continue

            if f_val is not None and str(f_val).strip() != "":
                try:
                    frame_idx = int(round(float(f_val)))
                except (ValueError, TypeError):
                    frame_idx = 0
            elif t_val is not None and str(t_val).strip() != "":
                try:
                    frame_idx = int(round(float(t_val) * 24))
                except (ValueError, TypeError):
                    frame_idx = 0
            else:
                frame_idx = 0

            guide_dict = {"frame_idx": max(0, frame_idx)}
            if saved_img:
                guide_dict["image"] = saved_img
            if saved_vid:
                guide_dict["video"] = saved_vid
            if saved_aud:
                guide_dict["audio"] = saved_aud
            active_guides.append(guide_dict)

    # 2. Check direct 'h3_guides' or 'guides' list of dicts (from MCP or custom payload)
    raw_guides = ui_values.get(f'h3_guides_{prefix}') or ui_values.get('h3_guides') or ui_values.get('guides', [])
    if raw_guides and not active_guides and isinstance(raw_guides, list):
        for item in raw_guides:
            if isinstance(item, dict):
                img = item.get('image')
                vid = item.get('video')
                aud = item.get('audio')

                saved_img = save_temp_image(img) if img is not None else None
                saved_vid = save_temp_video(vid) if vid else None
                saved_aud = save_temp_audio(aud) if aud else None

                if not (saved_img or saved_vid or saved_aud):
                    continue

                if item.get('frame_idx') is not None:
                    try:
                        frame_idx = int(round(float(item['frame_idx'])))
                    except (ValueError, TypeError):
                        frame_idx = 0
                elif item.get('time_seconds') is not None or item.get('time') is not None:
                    try:
                        raw_t = item.get('time_seconds') if item.get('time_seconds') is not None else item.get('time')
                        frame_idx = int(round(float(raw_t) * 24))
                    except (ValueError, TypeError):
                        frame_idx = 0
                else:
                    frame_idx = 0

                guide_dict = {"frame_idx": max(0, frame_idx)}
                if saved_img:
                    guide_dict["image"] = saved_img
                if saved_vid:
                    guide_dict["video"] = saved_vid
                if saved_aud:
                    guide_dict["audio"] = saved_aud
                active_guides.append(guide_dict)

    active_guides.sort(key=lambda x: x['frame_idx'])
    return active_guides