"""
Black Forest Labs Image Generation Configuration

Handles transformation between OpenAI-compatible format and Black Forest Labs API format
for image generation endpoints (flux-pro-1.1, flux-pro-1.1-ultra, flux-dev, flux-pro).

API Reference: https://docs.bfl.ai/
"""

import time
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union

import httpx

from litellm.llms.base_llm.image_generation.transformation import (
    BaseImageGenerationConfig,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
    AllMessageValues,
    OpenAIImageGenerationOptionalParams,
)
from litellm.types.utils import ImageObject, ImageResponse

from ..common_utils import (
    DEFAULT_API_BASE,
    IMAGE_GENERATION_MODELS,
    BlackForestLabsError,
)

if TYPE_CHECKING:
    from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj

    LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
    LiteLLMLoggingObj = Any


class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig):
    """
    Configuration for Black Forest Labs image generation (text-to-image).

    Supports:
    - flux-pro-1.1: Fast & reliable standard generation
    - flux-pro-1.1-ultra: Ultra high-resolution (up to 4MP)
    - flux-dev: Development/open-source variant
    - flux-pro: Original pro model

    Note: HTTP requests and polling are handled by the handler (handler.py).
    This class only handles data transformation.
    """

    def get_supported_openai_params(
        self, model: str
    ) -> List[OpenAIImageGenerationOptionalParams]:
        """
        Return list of OpenAI params supported by Black Forest Labs.

        Note: BFL uses different parameter names, these are mapped in map_openai_params.
        """
        return [
            "n",  # Number of images (BFL returns 1 per request, but ultra supports up to 4)
            "size",  # Maps to width/height or aspect_ratio
            "quality",  # Maps to raw mode for ultra
            "seed",
            "output_format",
            "safety_tolerance",
            "prompt_upsampling",
            "raw",
            "num_images",
            "image_url",
            "image_prompt_strength",
            "aspect_ratio",
        ]

    def map_openai_params(
        self,
        non_default_params: dict,
        optional_params: dict,
        model: str,
        drop_params: bool,
    ) -> dict:
        """
        Map OpenAI parameters to Black Forest Labs parameters.

        BFL-specific params are passed through directly.
        """
        supported_params = self.get_supported_openai_params(model)

        for k, v in non_default_params.items():
            if k in optional_params:
                continue

            if k in supported_params:
                # Map OpenAI 'size' to BFL width/height
                if k == "size" and v:
                    self._map_size_param(v, optional_params)
                elif k == "n":
                    if "ultra" in model.lower():
                        optional_params["num_images"] = v
                    # non-ultra: silently skip (n=1 is BFL default)
                elif k == "quality":
                    if v == "hd" and "ultra" in model.lower():
                        optional_params["raw"] = True
                    # other quality values have no BFL mapping
                else:
                    optional_params[k] = v
            elif not drop_params:
                raise ValueError(
                    f"Parameter {k} is not supported for model {model}. "
                    f"Supported parameters are {supported_params}. "
                    f"Set drop_params=True to drop unsupported parameters."
                )

        return optional_params

    def _map_size_param(self, size: str, optional_params: dict) -> None:
        """Map OpenAI size parameter to BFL width/height."""
        # Common size mappings
        size_mapping = {
            "1024x1024": (1024, 1024),
            "1792x1024": (1792, 1024),
            "1024x1792": (1024, 1792),
            "512x512": (512, 512),
            "256x256": (256, 256),
        }

        if size in size_mapping:
            width, height = size_mapping[size]
            optional_params["width"] = width
            optional_params["height"] = height
        elif "x" in size:
            # Parse custom size
            try:
                width, height = map(int, size.lower().split("x"))
                optional_params["width"] = width
                optional_params["height"] = height
            except ValueError:
                raise ValueError(
                    f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024')."
                )

    def validate_environment(
        self,
        headers: dict,
        model: str,
        messages: List[AllMessageValues],
        optional_params: dict,
        litellm_params: dict,
        api_key: Optional[str] = None,
        api_base: Optional[str] = None,
    ) -> dict:
        """
        Validate environment and set up headers for Black Forest Labs.

        BFL uses x-key header for authentication.
        """
        final_api_key: Optional[str] = (
            api_key
            or get_secret_str("BFL_API_KEY")
            or get_secret_str("BLACK_FOREST_LABS_API_KEY")
        )

        if not final_api_key:
            raise BlackForestLabsError(
                status_code=401,
                message="BFL_API_KEY is not set. Please set it via environment variable or pass api_key parameter.",
            )

        headers["x-key"] = final_api_key
        headers["Content-Type"] = "application/json"
        headers["Accept"] = "application/json"

        return headers

    def _get_model_endpoint(self, model: str) -> str:
        """
        Get the API endpoint for a given model.
        """
        # Remove provider prefix if present (e.g., "black_forest_labs/flux-pro-1.1")
        model_name = model.lower()
        if "/" in model_name:
            model_name = model_name.split("/")[-1]

        # Check if model is in our mapping
        if model_name in IMAGE_GENERATION_MODELS:
            return IMAGE_GENERATION_MODELS[model_name]

        raise ValueError(
            f"Unknown BFL image generation model: {model_name}. "
            f"Supported models: {list(IMAGE_GENERATION_MODELS.keys())}"
        )

    def get_complete_url(
        self,
        api_base: Optional[str],
        api_key: Optional[str],
        model: str,
        optional_params: dict,
        litellm_params: dict,
        stream: Optional[bool] = None,
    ) -> str:
        """
        Get the complete URL for the Black Forest Labs API request.
        """
        base_url: str = api_base or get_secret_str("BFL_API_BASE") or DEFAULT_API_BASE
        base_url = base_url.rstrip("/")

        endpoint = self._get_model_endpoint(model)
        return f"{base_url}{endpoint}"

    def transform_image_generation_request(
        self,
        model: str,
        prompt: str,
        optional_params: dict,
        litellm_params: dict,
        headers: dict,
    ) -> dict:
        """
        Transform OpenAI-style request to Black Forest Labs request format.

        https://docs.bfl.ai/flux_models/flux_1_1_pro
        """
        # Build request body with prompt
        request_body: Dict[str, Any] = {
            "prompt": prompt,
        }

        # BFL-specific params that can be passed through
        bfl_params = [
            "width",
            "height",
            "aspect_ratio",
            "seed",
            "output_format",
            "safety_tolerance",
            "prompt_upsampling",
            # Ultra-specific
            "raw",
            "num_images",
            "image_url",
            "image_prompt_strength",
        ]

        for param in bfl_params:
            if param in optional_params and optional_params[param] is not None:
                request_body[param] = optional_params[param]

        # Set default output format if not specified
        if "output_format" not in request_body:
            request_body["output_format"] = "png"

        return request_body

    def transform_image_generation_response(
        self,
        model: str,
        raw_response: httpx.Response,
        model_response: ImageResponse,
        logging_obj: LiteLLMLoggingObj,
        request_data: dict,
        optional_params: dict,
        litellm_params: dict,
        encoding: Any,
        api_key: Optional[str] = None,
        json_mode: Optional[bool] = None,
    ) -> ImageResponse:
        """
        Transform Black Forest Labs response to OpenAI-compatible ImageResponse.

        This is called with the FINAL polled response (after handler does polling).
        The response contains: {"status": "Ready", "result": {"sample": "https://..."}}
        """
        try:
            response_data = raw_response.json()
        except Exception as e:
            raise BlackForestLabsError(
                status_code=raw_response.status_code,
                message=f"Error parsing BFL response: {e}",
            )

        result = response_data.get("result", {})

        if not model_response.data:
            model_response.data = []

        # Handle single image (sample) or multiple images
        if isinstance(result, dict) and "sample" in result:
            model_response.data.append(ImageObject(url=result["sample"]))
        elif isinstance(result, list):
            # Multiple images returned
            for img in result:
                if isinstance(img, str):
                    model_response.data.append(ImageObject(url=img))
                elif isinstance(img, dict) and "url" in img:
                    model_response.data.append(ImageObject(url=img["url"]))

        if not model_response.data:
            raise BlackForestLabsError(
                status_code=500,
                message="No image URL in BFL result",
            )

        model_response.created = int(time.time())
        return model_response

    def get_error_class(
        self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
    ) -> BlackForestLabsError:
        """Return the appropriate error class for Black Forest Labs."""
        return BlackForestLabsError(
            status_code=status_code,
            message=error_message,
        )


def get_black_forest_labs_image_generation_config(
    model: str,
) -> BlackForestLabsImageGenerationConfig:
    """
    Get the appropriate image generation config for a Black Forest Labs model.

    Currently returns a single config class, but can be extended
    for model-specific configurations if needed.
    """
    return BlackForestLabsImageGenerationConfig()
