# 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
import httpx
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")


async def _save_url_image_to_path(image_url: str, target_path: str) -> str:
    os.makedirs(os.path.dirname(target_path), exist_ok=True)
    try:
        async with httpx.AsyncClient(follow_redirects=True) as client:
            response = await client.get(image_url, timeout=30.0)
            response.raise_for_status()

            if not os.path.splitext(target_path)[1]:
                content_type = response.headers.get("content-type", "").lower()
                url_path = image_url.split("?")[0]
                _, url_ext = os.path.splitext(url_path)
                url_ext = url_ext.lower()
                if url_ext in {".jpg", ".jpeg", ".png", ".webp", ".gif", ".bmp"}:
                    ext = ".jpg" if url_ext == ".jpeg" else url_ext
                elif "png" in content_type:
                    ext = ".png"
                elif "webp" in content_type:
                    ext = ".webp"
                else:
                    ext = ".jpg"
                target_path = f"{target_path}{ext}"

            async with aiofiles.open(target_path, "wb") as f:
                await f.write(response.content)
            return target_path
    except Exception as e:
        raise RuntimeError(f"Failed to download image from URL: {e}") from e


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