# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
import asyncio
import inspect
import json
import os
import shutil
import tempfile
import time
from contextlib import contextmanager
from typing import Any, Generator, List, Optional, Union

import httpx
from fastapi import HTTPException, UploadFile

from sglang.multimodal_gen.configs.sample.sampling_params import (
    DataType,
    SamplingParams,
)
from sglang.multimodal_gen.runtime.entrypoints.utils import (
    ListLorasReq,
    MergeLoraWeightsReq,
    SetLoraReq,
    ShutdownReq,
    UnmergeLoraWeightsReq,
    format_lora_message,
    save_outputs,
)
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
from sglang.multimodal_gen.runtime.scheduler_client import AsyncSchedulerClient
from sglang.multimodal_gen.runtime.server_args import get_global_server_args
from sglang.multimodal_gen.runtime.utils.common import parse_size
from sglang.multimodal_gen.runtime.utils.image_io import save_base64_image_to_path
from sglang.multimodal_gen.runtime.utils.logging_utils import (
    init_logger,
    log_batch_completion,
    log_generation_timer,
)
from sglang.multimodal_gen.runtime.utils.trace_wrapper import trace_req

# re-export LoRA protocol types for backward compatibility
__all__ = [
    "SetLoraReq",
    "MergeLoraWeightsReq",
    "UnmergeLoraWeightsReq",
    "ListLorasReq",
    "ShutdownReq",
    "format_lora_message",
]

logger = init_logger(__name__)

OUTPUT_QUALITY_MAPPER = {"maximum": 100, "high": 90, "medium": 55, "low": 35}
DEFAULT_FPS = 24
DEFAULT_VIDEO_SECONDS = 4


def _bad_request(message: str) -> HTTPException:
    return HTTPException(status_code=400, detail=message)


def _parse_size_or_raise(size: str) -> tuple[int, int]:
    width, height = parse_size(size)
    if width is None or height is None or width <= 0 or height <= 0:
        raise _bad_request("size must be formatted as positive WIDTHxHEIGHT")
    return width, height


def _validate_positive_int(kwargs: dict[str, Any], name: str) -> None:
    value = kwargs.get(name)
    if value is not None and int(value) <= 0:
        raise _bad_request(f"{name} must be positive")


def flatten_extra_params(payload: Any) -> dict[str, Any]:
    """Promote vLLM-Omni-style extra_params into regular request fields."""
    if not isinstance(payload, dict):
        return {}

    extra_params = payload.pop("extra_params", None)
    if isinstance(extra_params, str):
        try:
            extra_params = json.loads(extra_params)
        except Exception:
            extra_params = None
    if not isinstance(extra_params, dict):
        if "guardrails" in payload:
            payload.setdefault("use_guardrails", payload["guardrails"])
        return payload

    for key, value in extra_params.items():
        payload.setdefault(key, value)
    if "guardrails" in extra_params:
        payload.setdefault("use_guardrails", extra_params["guardrails"])

    return payload


@contextmanager
def temp_dir_if_disabled(
    configured_path: str | None,
) -> Generator[str, None, None]:
    """Yield *configured_path* when it is set, otherwise create a temporary
    directory that is automatically removed when the context exits."""
    if configured_path is not None:
        os.makedirs(configured_path, exist_ok=True)
        yield configured_path
    else:
        tmp = tempfile.mkdtemp(prefix="sglang_")
        try:
            yield tmp
        finally:
            shutil.rmtree(tmp, ignore_errors=True)


def choose_output_image_ext(
    output_format: Optional[str], background: Optional[str]
) -> str:
    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"


def build_sampling_params(request_id: str, **kwargs) -> SamplingParams:
    """Build SamplingParams from request parameters.

    Handles size parsing, output_quality resolution, and None filtering before
    delegating to SamplingParams.from_user_sampling_params_args. Callers pass
    only the parameters they have; None values are stripped automatically so
    that SamplingParams defaults apply.
    """
    server_args = get_global_server_args()

    # pop HTTP-layer params that aren't SamplingParams fields
    output_quality = kwargs.pop("output_quality", None)

    has_explicit_compression = kwargs.get("output_compression") is not None

    # parse "WxH" size string if provided
    size = kwargs.pop("size", None)
    if size:
        w, h = _parse_size_or_raise(size)
        # treat None dimensions as unset so parsed size can fill them
        if kwargs.get("width") is None:
            kwargs["width"] = w
        if kwargs.get("height") is None:
            kwargs["height"] = h

    for name in (
        "width",
        "height",
        "num_frames",
        "num_inference_steps",
        "num_outputs_per_prompt",
    ):
        _validate_positive_int(kwargs, name)

    # filter out None values to let SamplingParams defaults apply
    kwargs = {k: v for k, v in kwargs.items() if v is not None}
    kwargs.setdefault("save_output", True)

    sampling_params = SamplingParams.from_user_sampling_params_args(
        model_path=server_args.model_path,
        server_args=server_args,
        request_id=request_id,
        **kwargs,
    )

    # resolve output_quality → output_compression with the correct data_type.
    # SamplingParams.__post_init__ may have resolved with the wrong data_type
    # (default VIDEO) before _adjust() set the correct one.
    if not has_explicit_compression and output_quality is not None:
        resolved = adjust_output_quality(output_quality, sampling_params.data_type)
        if resolved is not None:
            sampling_params.output_compression = resolved

    return sampling_params


async def save_image_to_path(
    image: Union[UploadFile, bytes, str],
    target_path: str,
    *,
    prefer_remote_source: bool = False,
) -> str:
    input_path = await _maybe_url_image(
        image, target_path, prefer_remote_source=prefer_remote_source
    )
    if input_path is None:
        input_path = await _save_upload_to_path(image, target_path)
    return input_path


# Helpers
async def _save_upload_to_path(
    upload: Union[UploadFile, bytes], target_path: str
) -> str:
    os.makedirs(os.path.dirname(target_path), exist_ok=True)
    if isinstance(upload, bytes):
        content = upload
    elif isinstance(upload, (bytearray, memoryview)):
        content = bytes(upload)
    else:
        read = getattr(upload, "read", None)
        if not callable(read):
            raise TypeError(f"Unsupported image upload type: {type(upload).__name__}")
        content = read()
        if inspect.isawaitable(content):
            content = await content
        if isinstance(content, (bytearray, memoryview)):
            content = bytes(content)
        if not isinstance(content, bytes):
            raise TypeError(
                f"Image upload read() returned {type(content).__name__}, expected bytes"
            )
    with open(target_path, "wb") as f:
        f.write(content)
    return target_path


async def _maybe_url_image(
    img_url: str,
    target_path: str,
    *,
    prefer_remote_source: bool = False,
) -> str | None:
    if not isinstance(img_url, str):
        return None

    if img_url.lower().startswith(("http://", "https://")):
        # Only bypass persistence when the caller explicitly disables input saves.
        # Otherwise keep the prefetch outside the measured server stages.
        if prefer_remote_source:
            return img_url
        # download image from URL and persist on disk
        input_path = await _save_url_image_to_path(img_url, target_path)
        return input_path
    elif img_url.startswith("data:image"):
        if prefer_remote_source:
            return img_url
        # encode image base64 url and persist on disk
        input_path = save_base64_image_to_path(img_url, target_path)
        return input_path
    else:
        raise ValueError("Unsupported image url format")


async def _save_url_image_to_path(image_url: str, target_path: str) -> str:
    """Download image from URL and save to target path."""

    def _is_retryable_download_error(error: Exception) -> bool:
        if isinstance(error, httpx.HTTPStatusError):
            status_code = error.response.status_code
            # Retry on rate limit and transient server-side failures.
            return status_code == 429 or 500 <= status_code < 600
        # Retry on transient network/protocol issues.
        return isinstance(
            error,
            (
                httpx.TimeoutException,
                httpx.NetworkError,
                httpx.RemoteProtocolError,
            ),
        )

    os.makedirs(os.path.dirname(target_path), exist_ok=True)

    max_attempts = 3
    backoff_seconds = 0.2
    last_error: Exception | None = None

    try:
        async with httpx.AsyncClient(follow_redirects=True) as client:
            for attempt in range(1, max_attempts + 1):
                try:
                    response = await client.get(image_url, timeout=10.0)
                    response.raise_for_status()

                    # Determine file extension from content type or URL after downloading
                    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 content_type.startswith("image/"):
                            if "jpeg" in content_type or "jpg" in content_type:
                                ext = ".jpg"
                            elif "png" in content_type:
                                ext = ".png"
                            elif "webp" in content_type:
                                ext = ".webp"
                            else:
                                ext = ".jpg"  # Default to jpg
                        elif content_type == "application/octet-stream":
                            # for octet-stream, if we couldn't get it from URL, default to jpg
                            ext = ".jpg"
                        else:
                            raise ValueError(
                                f"URL does not point to an image. Content-Type: {content_type}"
                            )
                        target_path = f"{target_path}{ext}"

                    with open(target_path, "wb") as f:
                        f.write(response.content)

                    return target_path
                except Exception as e:
                    last_error = e
                    if attempt == max_attempts or not _is_retryable_download_error(e):
                        raise
                    wait_s = backoff_seconds * (2 ** (attempt - 1))
                    logger.warning(
                        "Retrying image download (%s/%s) for %s after %.1fs due to: %s",
                        attempt,
                        max_attempts,
                        image_url,
                        wait_s,
                        e,
                    )
                    await asyncio.sleep(wait_s)
    except Exception as e:
        final_error = last_error or e
        raise Exception(
            f"Failed to download image from URL {image_url}: {str(final_error)}"
        )


async def process_generation_batch(
    scheduler_client: AsyncSchedulerClient,
    batch,
) -> tuple[list[str], OutputBatch]:
    total_start_time = time.perf_counter()
    with trace_req(batch.trace_ctx), log_generation_timer(logger, batch.prompt):
        result = await scheduler_client.forward([batch])

        if (
            result.output is None
            and result.output_file_paths is None
            and result.raw_frame_batches is None
        ):
            error_msg = result.error or "Unknown error"
            raise RuntimeError(
                f"Model generation returned no output. Error from scheduler: {error_msg}"
            )

        save_file_path_list = []
        if result.output_file_paths:
            save_file_path_list = result.output_file_paths
        elif result.output is not None:
            num_outputs = len(result.output)
            save_file_path_list = save_outputs(
                result.output,
                batch.data_type,
                batch.fps,
                batch.save_output,
                lambda idx: str(batch.output_file_path(num_outputs, idx)),
                audio=result.audio,
                audio_sample_rate=result.audio_sample_rate,
                output_compression=batch.output_compression,
                enable_frame_interpolation=batch.enable_frame_interpolation,
                frame_interpolation_exp=batch.frame_interpolation_exp,
                frame_interpolation_scale=batch.frame_interpolation_scale,
                frame_interpolation_model_path=batch.frame_interpolation_model_path,
                enable_upscaling=batch.enable_upscaling,
                upscaling_model_path=batch.upscaling_model_path,
                upscaling_scale=batch.upscaling_scale,
            )

    total_time = time.perf_counter() - total_start_time
    if get_global_server_args().batching_max_size > 1:
        log_batch_completion(
            logger,
            len(save_file_path_list),
            total_time,
        )

    if result.peak_memory_mb and result.peak_memory_mb > 0:
        logger.info(f"Peak memory usage: {result.peak_memory_mb:.2f} MB")

    return save_file_path_list, result


def merge_image_input_list(*inputs: Union[List, Any, None]) -> List:
    """
    Merge multiple image input sources into a single list.

    This function handles both single items and lists of items, merging them
    into a single flattened list. Useful for processing images, URLs, or other
    multimedia inputs that can come as either single items or lists.

    Args:
        *inputs: Variable number of inputs, each can be None, single item, or list

    Returns:
        List: Flattened list of all non-None inputs

    Example:
        >>> merge_image_input_list(["img1", "img2"], "img3", None)
        ["img1", "img2", "img3"]
    """
    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


def add_common_data_to_response(
    response: dict, request_id: str, result: OutputBatch
) -> dict:
    if result.peak_memory_mb and result.peak_memory_mb > 0:
        response["peak_memory_mb"] = result.peak_memory_mb

    if result.metrics and result.metrics.total_duration_s > 0:
        response["inference_time_s"] = result.metrics.total_duration_s

    response["id"] = request_id

    if result.action_pred is not None:
        t = result.action_pred
        response["action"] = {
            "data": t.tolist(),
            "shape": list(t.shape),
            "dtype": str(t.dtype).replace("torch.", ""),
            "raw_action_dim": result.action_raw_action_dim,
            "action_mode": result.action_mode,
            "domain_id": result.action_domain_id,
        }

    return response


def adjust_output_quality(output_quality: str, data_type: DataType = None) -> int:
    if output_quality == "default":
        return 50 if data_type == DataType.VIDEO else 75
    return OUTPUT_QUALITY_MAPPER.get(output_quality, None)