# Adapted from SGLang
# (https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py)

import base64
import os
import re
from typing import Any

import aiofiles
from fastapi import UploadFile

from fastvideo.logger import init_logger

logger = init_logger(__name__)


def parse_size(size: str) -> tuple[int, int] | tuple[None, None]:
    """Parse a 'WIDTHxHEIGHT' string into (width, height)"""
    try:
        parts = size.lower().replace(" ", "").split("x")
        if len(parts) != 2:
            raise ValueError
        w, h = int(parts[0]), int(parts[1])
        return w, h
    except Exception:
        return None, None


def choose_image_ext(output_format: str | None, background: str | None) -> str:
    """Pick a file extension for image outputs"""
    fmt = (output_format or "").lower()
    if fmt in {"png", "webp", "jpeg", "jpg"}:
        return "jpg" if fmt == "jpeg" else fmt
    if (background or "auto").lower() == "transparent":
        return "png"
    return "jpg"


async def save_image_to_path(image: UploadFile | str, target_path: str) -> str:
    """Save an uploaded file or download from URL to *target_path*"""
    input_path = await _maybe_url_image(image, target_path)
    if input_path is None:
        input_path = await _save_upload_to_path(image, target_path)
    return input_path


async def _save_upload_to_path(upload: UploadFile, target_path: str) -> str:
    os.makedirs(os.path.dirname(target_path), exist_ok=True)
    content = await upload.read()
    async with aiofiles.open(target_path, "wb") as f:
        await f.write(content)
    return target_path


async def _maybe_url_image(img_url: str, target_path: str) -> str | None:
    if not isinstance(img_url, str):
        return None
    if img_url.lower().startswith(("http://", "https://")):
        return await _save_url_image_to_path(img_url, target_path)
    elif img_url.startswith("data:image"):
        return await save_base64_image_to_path(img_url, target_path)
    else:
        raise ValueError("Unsupported image url format")


_IMAGE_SUFFIXES = frozenset({".jpg", ".jpeg", ".png", ".webp", ".gif", ".bmp"})


def image_suffix(image_url: str) -> str:
    """File suffix for a downloaded image; decoding inspects content, so this is only a hint."""
    suffix = os.path.splitext(image_url.split("?")[0])[1].lower()
    if suffix not in _IMAGE_SUFFIXES:
        return ".jpg"
    return ".jpg" if suffix == ".jpeg" else suffix


async def _save_url_image_to_path(image_url: str, target_path: str) -> str:
    # Local imports: state imports serving modules that import these helpers.
    from fastvideo.entrypoints.openai.media_policy import fetch_media
    from fastvideo.entrypoints.openai.state import get_media_policy, get_media_reference_max_bytes

    if not os.path.splitext(target_path)[1]:
        target_path = f"{target_path}{image_suffix(image_url)}"
    try:
        await fetch_media(image_url, target_path, max_bytes=get_media_reference_max_bytes(), policy=get_media_policy())
    except Exception as e:
        raise RuntimeError(f"Failed to download image from URL: {e}") from e
    return target_path


async def save_base64_image_to_path(base64_data: str, target_path: str) -> str:
    pattern = r"data:(.*?)(;base64)?,(.*)"
    match = re.match(pattern, base64_data)
    if not match or not match.group(2):
        raise ValueError("Invalid base64 image data URL")
    media_type = match.group(1)
    data = match.group(3)
    if not data:
        raise ValueError("Empty base64 image data")
    ext = media_type.split("/")[-1].lower() if media_type.startswith("image/") else "jpg"
    if ext == "jpeg":
        ext = "jpg"
    target_path = f"{target_path}.{ext}"
    os.makedirs(os.path.dirname(target_path), exist_ok=True)
    try:
        image_data = base64.b64decode(data)
        async with aiofiles.open(target_path, "wb") as f:
            await f.write(image_data)
        return target_path
    except Exception as e:
        raise RuntimeError(f"Failed to decode base64 image: {e}") from e


def merge_image_input_list(*inputs: list | Any | None) -> list:
    """Merge multiple image input sources into a single flat list"""
    result = []
    for input_item in inputs:
        if input_item is not None:
            if isinstance(input_item, list):
                result.extend(input_item)
            else:
                result.append(input_item)
    return result