"""
MCP Common Utilities & Data Structures
Contains image parsing utilities, task definitions, aspect ratio presets, and async task database for MiniMax H3 Video Generation.
"""

import os
import time
import uuid
import urllib.parse
import urllib.request
import urllib.error
import base64
import io
import yaml
from typing import Dict, Any
from PIL import Image

_PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
_YAML_DIR = os.path.join(_PROJECT_ROOT, "yaml")

_MAX_IMAGE_DOWNLOAD_BYTES = 50 * 1024 * 1024  # 50 MB
_IMAGE_DOWNLOAD_TIMEOUT = 30  # seconds
_MAX_VIDEO_DOWNLOAD_BYTES = 200 * 1024 * 1024  # 200 MB
_VIDEO_DOWNLOAD_TIMEOUT = 60  # seconds
_MAX_AUDIO_DOWNLOAD_BYTES = 100 * 1024 * 1024  # 100 MB
_AUDIO_DOWNLOAD_TIMEOUT = 60  # seconds

_ALLOWED_IMAGE_CONTENT_TYPES = frozenset([
    "image/png", "image/jpeg", "image/jpg", "image/gif",
    "image/webp", "image/bmp", "image/tiff",
])

ASPECT_RATIO_PRESETS = {
    "16:9 (Landscape)": {"width": 1344, "height": 768},
    "9:16 (Portrait)": {"width": 768, "height": 1344},
    "1:1 (Square)": {"width": 1024, "height": 1024},
    "4:3 (Classic TV)": {"width": 1152, "height": 864},
    "3:4 (Classic Portrait)": {"width": 864, "height": 1152},
    "3:2 (Landscape)": {"width": 1248, "height": 832},
    "2:3 (Portrait)": {"width": 832, "height": 1248},
}


def _parse_image_param(image_param: Any) -> Any:
    """Parse a Base64 Data URI, HTTP/HTTPS URL, or PIL.Image into a PIL Image object."""
    if isinstance(image_param, Image.Image):
        return image_param

    if not isinstance(image_param, str) or not image_param.strip():
        return None

    image_param = image_param.strip()

    if os.path.exists(image_param):
        return Image.open(image_param)

    if image_param.startswith("http://") or image_param.startswith("https://"):
        return _download_image_from_url(image_param)

    if image_param.startswith("data:image/"):
        _, encoded = image_param.split(",", 1) if "," in image_param else ("", image_param)
        data = base64.b64decode(encoded)
        return Image.open(io.BytesIO(data))

    if len(image_param) > 100:
        try:
            data = base64.b64decode(image_param)
            return Image.open(io.BytesIO(data))
        except Exception:
            pass

    raise ValueError(
        "Invalid image parameter format. Expected a file path, Base64 Data URI, or an HTTP/HTTPS URL."
    )


def _download_image_from_url(url: str) -> Image.Image:
    """Download an image from an HTTP/HTTPS URL and return it as a PIL Image."""
    req = urllib.request.Request(url, headers={"User-Agent": "H3-VideoGen-MCP/1.0"})
    try:
        with urllib.request.urlopen(req, timeout=_IMAGE_DOWNLOAD_TIMEOUT) as resp:
            content_type = resp.headers.get("Content-Type", "").split(";")[0].strip().lower()
            if content_type and content_type not in _ALLOWED_IMAGE_CONTENT_TYPES:
                raise ValueError(
                    f"URL returned non-image Content-Type '{content_type}'. "
                    f"Expected one of: {', '.join(sorted(_ALLOWED_IMAGE_CONTENT_TYPES))}."
                )

            content_length = resp.headers.get("Content-Length")
            if content_length and int(content_length) > _MAX_IMAGE_DOWNLOAD_BYTES:
                raise ValueError(
                    f"Image at URL is too large ({int(content_length)} bytes). "
                    f"Maximum allowed size is {_MAX_IMAGE_DOWNLOAD_BYTES} bytes."
                )

            chunks = []
            total = 0
            while True:
                chunk = resp.read(8192)
                if not chunk:
                    break
                total += len(chunk)
                if total > _MAX_IMAGE_DOWNLOAD_BYTES:
                    raise ValueError(f"Image download exceeded maximum allowed size of {_MAX_IMAGE_DOWNLOAD_BYTES} bytes.")
                chunks.append(chunk)

            data = b"".join(chunks)

    except urllib.error.URLError as e:
        raise ValueError(f"Failed to download image from URL: {e}") from e
    except urllib.error.HTTPError as e:
        raise ValueError(f"HTTP error {e.code} when downloading image from URL: {e.reason}") from e

    if not data:
        raise ValueError("Downloaded image data is empty.")

    return Image.open(io.BytesIO(data))


def _download_file_from_url(url: str, suffix: str = ".mp4", max_bytes: int = _MAX_VIDEO_DOWNLOAD_BYTES, timeout: int = _VIDEO_DOWNLOAD_TIMEOUT) -> str:
    """Download a media file from an HTTP/HTTPS URL and save it to a temporary local path."""
    req = urllib.request.Request(url, headers={"User-Agent": "H3-VideoGen-MCP/1.0"})
    input_dir = os.path.join(_PROJECT_ROOT, "input")
    os.makedirs(input_dir, exist_ok=True)
    temp_path = os.path.join(input_dir, f"mcp_download_{uuid.uuid4().hex[:10]}{suffix}")
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            content_length = resp.headers.get("Content-Length")
            if content_length and int(content_length) > max_bytes:
                raise ValueError(
                    f"File at URL is too large ({int(content_length)} bytes). "
                    f"Maximum allowed size is {max_bytes} bytes."
                )

            total = 0
            with open(temp_path, "wb") as f:
                while True:
                    chunk = resp.read(65536)
                    if not chunk:
                        break
                    total += len(chunk)
                    if total > max_bytes:
                        raise ValueError(f"File download exceeded maximum allowed size of {max_bytes} bytes.")
                    f.write(chunk)
        return temp_path
    except Exception as e:
        if os.path.exists(temp_path):
            try:
                os.remove(temp_path)
            except Exception:
                pass
        raise ValueError(f"Failed to download file from URL '{url}': {e}") from e


def _parse_video_param(video_param: Any) -> Any:
    """Parse a file path, Base64 Data URI, or an HTTP/HTTPS URL into a local video file path."""
    if not isinstance(video_param, str) or not video_param.strip():
        return None

    video_param = video_param.strip()

    if os.path.exists(video_param):
        return video_param

    if video_param.startswith("http://") or video_param.startswith("https://"):
        parsed_url = urllib.parse.urlparse(video_param)
        ext = os.path.splitext(parsed_url.path)[1] or ".mp4"
        return _download_file_from_url(video_param, suffix=ext, max_bytes=_MAX_VIDEO_DOWNLOAD_BYTES, timeout=_VIDEO_DOWNLOAD_TIMEOUT)

    if video_param.startswith("data:video/"):
        _, encoded = video_param.split(",", 1) if "," in video_param else ("", video_param)
        data = base64.b64decode(encoded)
        input_dir = os.path.join(_PROJECT_ROOT, "input")
        os.makedirs(input_dir, exist_ok=True)
        temp_path = os.path.join(input_dir, f"mcp_upload_{uuid.uuid4().hex[:10]}.mp4")
        with open(temp_path, "wb") as f:
            f.write(data)
        return temp_path

    return video_param


def _parse_audio_param(audio_param: Any) -> Any:
    """Parse a file path, Base64 Data URI, or an HTTP/HTTPS URL into a local audio file path."""
    if not isinstance(audio_param, str) or not audio_param.strip():
        return None

    audio_param = audio_param.strip()

    if os.path.exists(audio_param):
        return audio_param

    if audio_param.startswith("http://") or audio_param.startswith("https://"):
        parsed_url = urllib.parse.urlparse(audio_param)
        ext = os.path.splitext(parsed_url.path)[1] or ".mp3"
        return _download_file_from_url(audio_param, suffix=ext, max_bytes=_MAX_AUDIO_DOWNLOAD_BYTES, timeout=_AUDIO_DOWNLOAD_TIMEOUT)

    if audio_param.startswith("data:audio/"):
        _, encoded = audio_param.split(",", 1) if "," in audio_param else ("", audio_param)
        data = base64.b64decode(encoded)
        input_dir = os.path.join(_PROJECT_ROOT, "input")
        os.makedirs(input_dir, exist_ok=True)
        temp_path = os.path.join(input_dir, f"mcp_upload_{uuid.uuid4().hex[:10]}.mp3")
        with open(temp_path, "wb") as f:
            f.write(data)
        return temp_path

    return audio_param


def _load_yaml(filepath: str) -> dict:
    """Safely load a YAML file."""
    if not os.path.exists(filepath):
        print(f"Warning: YAML file not found: {filepath}")
        return {}
    with open(filepath, "r", encoding="utf-8") as f:
        return yaml.safe_load(f) or {}


_GLOBAL_OPTIONAL_INPUTS_SCHEMA = {
    "seed": {
        "type": "integer",
        "default": -1,
        "description": "Random seed for video generation (-1 for random seed, >=0 for deterministic reproduction)."
    },
    "zero_gpu_duration": {
        "type": "integer",
        "default": 120,
        "minimum": 1,
        "description": "GPU execution time quota allocation in seconds on HuggingFace ZeroGPU spaces (default: 120)."
    },
    "loras": {
        "type": "array",
        "description": "List of LoRA configurations to apply. E.g., [{'source': 'Hugging Face', 'id_or_url': 'repo/lora.safetensors', 'scale': 1.0}] or [{'lora_name': 'my_lora.safetensors', 'strength_model': 1.0, 'strength_clip': 1.0}].",
        "items": {
            "type": "object"
        }
    },
    "h3_controlnets": {
        "type": "array",
        "description": "List of H3 ControlNet configurations. E.g., [{'video': 'https://example.com/pose.mp4', 'control_net_name': 'minimax_h3_fun_controlnet_union_pruned_int8_convrot.safetensors', 'strength': 1.0, 'start_percent': 0.0, 'end_percent': 1.0}].",
        "items": {
            "type": "object"
        }
    },
    "control_video": {
        "type": "string",
        "description": "Convenience parameter: URL, file path, or Base64 Data URI of a single control video for H3 ControlNet."
    },
    "control_strength": {
        "type": "number",
        "default": 1.0,
        "description": "Strength for single convenience control video (default: 1.0)."
    },
    "control_net_name": {
        "type": "string",
        "default": "minimax_h3_fun_controlnet_union_pruned_int8_convrot.safetensors",
        "description": "Model filename for H3 ControlNet (default: minimax_h3_fun_controlnet_union_pruned_int8_convrot.safetensors)."
    },
    "h3_guides": {
        "type": "array",
        "description": "List of MiniMax H3 timeline Keyframe Guide configurations (Add Guide for MiniMax H3). E.g., [{'image': 'https://example.com/frame.png', 'video': 'https://example.com/clip.mp4', 'audio': 'https://example.com/voice.mp3', 'frame_idx': 36, 'time_seconds': 1.5}].",
        "items": {
            "type": "object"
        }
    }
}

_COMMON_OPTIONAL_INPUTS = [
    "seed", "zero_gpu_duration", "loras", "h3_controlnets", "control_video", "control_strength", "control_net_name", "h3_guides"
]

_TASK_DEFINITIONS = [
    {
        "task_type": "t2va",
        "display_name": "Text-to-Video (H3 FL2VA)",
        "description": "Generate videos from text prompt without input frames. Supports LoRAs and H3 ControlNet.",
        "required_inputs": ["prompt", "width", "height", "duration"],
        "optional_inputs": _COMMON_OPTIONAL_INPUTS,
        "optional_inputs_schema": _GLOBAL_OPTIONAL_INPUTS_SCHEMA,
        "aspect_ratio_presets": ASPECT_RATIO_PRESETS,
        "example_json_params": {
            "task_type": "t2va",
            "prompt": "A futuristic city with flying cars at sunset, cinematic 4k",
            "width": 1344,
            "height": 768,
            "duration": 5.0,
            "seed": -1,
            "zero_gpu_duration": 120
        }
    },
    {
        "task_type": "i2va",
        "display_name": "Image-to-Video (H3 FL2VA)",
        "description": "Generate videos starting from a single initial image frame (first frame). Supports LoRAs and H3 ControlNet.",
        "required_inputs": ["prompt", "width", "height", "duration", "first_frame_image"],
        "optional_inputs": _COMMON_OPTIONAL_INPUTS,
        "optional_inputs_schema": _GLOBAL_OPTIONAL_INPUTS_SCHEMA,
        "aspect_ratio_presets": ASPECT_RATIO_PRESETS,
        "example_json_params": {
            "task_type": "i2va",
            "prompt": "A majestic lion standing on a rock looking at the horizon",
            "width": 1344,
            "height": 768,
            "duration": 5.0,
            "first_frame_image": "https://example.com/first_frame.png",
            "seed": -1,
            "zero_gpu_duration": 120
        }
    },
    {
        "task_type": "flf2va",
        "display_name": "First & Last Frame-to-Video (H3 FL2VA)",
        "description": "Generate videos guided by both a starting frame (first frame) and ending frame (last frame). Supports LoRAs and H3 ControlNet.",
        "required_inputs": ["prompt", "width", "height", "duration", "first_frame_image", "last_frame_image"],
        "optional_inputs": _COMMON_OPTIONAL_INPUTS,
        "optional_inputs_schema": _GLOBAL_OPTIONAL_INPUTS_SCHEMA,
        "aspect_ratio_presets": ASPECT_RATIO_PRESETS,
        "example_json_params": {
            "task_type": "flf2va",
            "prompt": "Smooth transition of a flower blooming from start to finish",
            "width": 1344,
            "height": 768,
            "duration": 5.0,
            "first_frame_image": "https://example.com/first_frame.png",
            "last_frame_image": "https://example.com/last_frame.png",
            "seed": -1,
            "zero_gpu_duration": 120
        }
    },
    {
        "task_type": "ref2va",
        "display_name": "Reference-to-Video (H3 REF2VA)",
        "description": "Generate videos with multi-modal references (up to 9 ref images, 3 ref videos, and 3 ref audios). Supports LoRAs, H3 ControlNet, and timeline Keyframe Guides (Add Guide for MiniMax H3).",
        "required_inputs": ["prompt", "width", "height", "duration"],
        "optional_inputs": [
            *_COMMON_OPTIONAL_INPUTS,
            "ref_image1", "ref_image2", "ref_image3", "ref_image4", "ref_image5", "ref_image6", "ref_image7", "ref_image8", "ref_image9",
            "ref_video1", "ref_video2", "ref_video3",
            "ref_audio1", "ref_audio2", "ref_audio3",
            "h3_guides"
        ],
        "optional_inputs_schema": _GLOBAL_OPTIONAL_INPUTS_SCHEMA,
        "aspect_ratio_presets": ASPECT_RATIO_PRESETS,
        "example_json_params": {
            "task_type": "ref2va",
            "prompt": "A character dancing in a cyberpunk room guided by reference media",
            "width": 1344,
            "height": 768,
            "duration": 5.0,
            "ref_image1": "https://example.com/ref_character.png",
            "ref_video1": "https://example.com/ref_motion.mp4",
            "ref_audio1": "https://example.com/ref_bgm.mp3",
            "seed": -1,
            "zero_gpu_duration": 120
        }
    }
]

_TASKS_DB: Dict[str, Dict[str, Any]] = {}


class DummyProgress:
    def __call__(self, progress=0.0, desc=None):
        pass


def _get_public_base_url() -> str:
    """Auto-resolve the publicly accessible base URL."""
    public_url = os.getenv("PUBLIC_URL") or os.getenv("BASE_URL")
    if public_url:
        return public_url.rstrip("/")

    space_host = os.getenv("SPACE_HOST")
    if space_host:
        if not space_host.startswith("http://") and not space_host.startswith("https://"):
            return f"https://{space_host}"
        return space_host.rstrip("/")

    try:
        from core.settings import GRADIO_SERVER_NAME, SERVER_PORT
    except ImportError:
        GRADIO_SERVER_NAME = "127.0.0.1"
        SERVER_PORT = 7860

    server_name = os.getenv("GRADIO_SERVER_NAME", GRADIO_SERVER_NAME)
    if server_name == "0.0.0.0":
        server_name = "127.0.0.1"
    port = os.getenv("GRADIO_SERVER_PORT", str(SERVER_PORT))

    return f"http://{server_name}:{port}"


def _execute_imagegen_pipeline(task_id: str, params: dict):
    """Execute the MiniMax H3 video generation pipeline in the background and update _TASKS_DB."""
    start_time = time.time()
    try:
        _TASKS_DB[task_id]["status"] = "processing"
        _TASKS_DB[task_id]["progress"] = 10
        _TASKS_DB[task_id]["updated_at"] = int(start_time)

        task_type = params.get("task_type", "t2va")
        prompt = params.get("prompt", "")
        width = int(params.get("width", 1344))
        height = int(params.get("height", 768))
        duration = float(params.get("duration", 5.0))
        seed = int(params.get("seed", -1))
        zero_gpu_duration = int(params.get("zero_gpu_duration", 120))

        # Process LoRAs if provided
        loras_param = params.get("loras")
        loras_list = loras_param if isinstance(loras_param, list) else []

        # Process H3 ControlNets if provided
        h3_cns_param = params.get("h3_controlnets")
        parsed_cns = []
        if h3_cns_param and isinstance(h3_cns_param, list):
            for item in h3_cns_param:
                if isinstance(item, dict):
                    vid_raw = item.get("video") or item.get("control_video") or item.get("file")
                    vid_parsed = _parse_video_param(vid_raw) if vid_raw else None
                    if vid_parsed:
                        new_item = dict(item)
                        new_item["video"] = vid_parsed
                        parsed_cns.append(new_item)
        elif params.get("control_video"):
            vid_raw = params.get("control_video")
            vid_parsed = _parse_video_param(vid_raw) if vid_raw else None
            if vid_parsed:
                parsed_cns.append({
                    "video": vid_parsed,
                    "control_net_name": params.get("control_net_name") or "minimax_h3_fun_controlnet_union_pruned_int8_convrot.safetensors",
                    "strength": float(params.get("control_strength", 1.0)),
                    "start_percent": float(params.get("control_start_percent", 0.0)),
                    "end_percent": float(params.get("control_end_percent", 1.0)),
                })

        parsed_guides = []
        raw_guides = params.get("h3_guides") or params.get("guides")
        if raw_guides and isinstance(raw_guides, list):
            for item in raw_guides:
                if isinstance(item, dict):
                    guide_entry = {}
                    if item.get("image"):
                        guide_entry["image"] = _parse_image_param(item["image"])
                    if item.get("video"):
                        guide_entry["video"] = _parse_video_param(item["video"])
                    if item.get("audio"):
                        guide_entry["audio"] = _parse_audio_param(item["audio"])
                    if item.get("frame_idx") is not None:
                        guide_entry["frame_idx"] = item["frame_idx"]
                    if item.get("time_seconds") is not None:
                        guide_entry["time_seconds"] = item["time_seconds"]
                    if guide_entry:
                        parsed_guides.append(guide_entry)

        if task_type in ("t2va", "i2va", "flf2va"):
            from ui.h3_logic import process_inputs as process_h3_inputs
            
            first_frame_raw = params.get("first_frame_image") or params.get("first_frame")
            last_frame_raw = params.get("last_frame_image") or params.get("last_frame")

            first_frame_pil = _parse_image_param(first_frame_raw) if first_frame_raw else None
            last_frame_pil = _parse_image_param(last_frame_raw) if last_frame_raw else None

            ui_inputs = {
                "task": task_type.upper(),
                "first_frame": first_frame_pil,
                "last_frame": last_frame_pil,
                "prompt": prompt,
                "resolution": "720p",
                "aspect_ratio": "Custom",
                "width": width,
                "height": height,
                "duration": duration,
                "seed": seed,
                "zero_gpu_duration": zero_gpu_duration,
                "batch_count": 1,
                "loras": loras_list,
                "h3_controlnets": parsed_cns,
                "h3_guides": parsed_guides,
            }
        elif task_type == "ref2va":
            from ui.h3_logic import process_inputs as process_h3_inputs

            ref_images = []
            for i in range(1, 10):
                img_val = params.get(f"ref_image{i}")
                if img_val:
                    pil_img = _parse_image_param(img_val)
                    if pil_img:
                        ref_images.append(pil_img)
            if not ref_images and "ref_images" in params and isinstance(params["ref_images"], list):
                for img_val in params["ref_images"]:
                    if img_val:
                        pil_img = _parse_image_param(img_val)
                        if pil_img:
                            ref_images.append(pil_img)

            ref_videos = []
            for i in range(1, 4):
                vid_val = params.get(f"ref_video{i}")
                if vid_val:
                    parsed_vid = _parse_video_param(vid_val)
                    if parsed_vid:
                        ref_videos.append(parsed_vid)
            if not ref_videos and "ref_videos" in params and isinstance(params["ref_videos"], list):
                for vid_val in params["ref_videos"]:
                    if vid_val:
                        parsed_vid = _parse_video_param(vid_val)
                        if parsed_vid:
                            ref_videos.append(parsed_vid)

            ref_audios = []
            for i in range(1, 4):
                aud_val = params.get(f"ref_audio{i}")
                if aud_val:
                    parsed_aud = _parse_audio_param(aud_val)
                    if parsed_aud:
                        ref_audios.append(parsed_aud)
            if not ref_audios and "ref_audios" in params and isinstance(params["ref_audios"], list):
                for aud_val in params["ref_audios"]:
                    if aud_val:
                        parsed_aud = _parse_audio_param(aud_val)
                        if parsed_aud:
                            ref_audios.append(parsed_aud)

            ui_inputs = {
                "task": "REF2VA",
                "prompt": prompt,
                "resolution": "720p",
                "aspect_ratio": "Custom",
                "width": width,
                "height": height,
                "duration": duration,
                "seed": seed,
                "zero_gpu_duration": zero_gpu_duration,
                "batch_count": 1,
                "ref_images": ref_images,
                "ref_videos": ref_videos,
                "ref_audios": ref_audios,
                "loras": loras_list,
                "h3_controlnets": parsed_cns,
                "h3_guides": parsed_guides,
            }
        else:
            raise ValueError(f"Unsupported task_type '{task_type}'")

        _TASKS_DB[task_id]["progress"] = 50

        from core.generation_logic import generate_h3_video_wrapper
        output = generate_h3_video_wrapper(process_h3_inputs, ui_inputs, progress=DummyProgress())

        try:
            from core.settings import OUTPUT_DIR
        except ImportError:
            OUTPUT_DIR = os.path.join(_PROJECT_ROOT, "output")

        os.makedirs(OUTPUT_DIR, exist_ok=True)

        import tempfile
        import gradio.processing_utils as pu

        gradio_cache_dir = os.path.join(tempfile.gettempdir(), "gradio")
        os.makedirs(gradio_cache_dir, exist_ok=True)

        base_url = _get_public_base_url()
        videos = []
        raw_list = output if isinstance(output, list) else ([output] if output else [])
        for idx, item in enumerate(raw_list):
            target_path = None
            if isinstance(item, str) and os.path.exists(item):
                target_path = item

            if target_path:
                try:
                    cached_path = pu.save_file_to_cache(target_path, cache_dir=gradio_cache_dir)
                    abs_path = os.path.abspath(cached_path).replace("\\", "/")
                except Exception as e:
                    print(f"Warning: Failed to cache video file to Gradio temp dir: {e}")
                    abs_path = os.path.abspath(target_path).replace("\\", "/")

                url = f"{base_url}/gradio_api/file={urllib.parse.quote(abs_path)}"
                videos.append(url)
            elif item:
                videos.append(str(item))

        execution_time = round(time.time() - start_time, 2)
        _TASKS_DB[task_id]["status"] = "completed"
        _TASKS_DB[task_id]["progress"] = 100
        _TASKS_DB[task_id]["completed_at"] = int(time.time())
        _TASKS_DB[task_id]["result"] = {
            "task_type": task_type,
            "videos": videos,
            "video_path": raw_list[0] if raw_list and isinstance(raw_list[0], str) else None,
            "seed": seed,
            "width": width,
            "height": height,
            "duration": duration,
            "execution_time_seconds": execution_time,
        }

    except Exception as e:
        _TASKS_DB[task_id]["status"] = "failed"
        _TASKS_DB[task_id]["progress"] = 0
        _TASKS_DB[task_id]["failed_at"] = int(time.time())
        _TASKS_DB[task_id]["error"] = {
            "code": "EXECUTION_ERROR",
            "message": str(e),
        }