"""
Black Forest Labs Image Generation Handler

Handles image generation requests for Black Forest Labs models.
BFL uses an async polling pattern - the initial request returns a task ID,
then we poll until the result is ready.
"""

import asyncio
import time
from collections.abc import Coroutine, Mapping
from typing import Final, Protocol, TypedDict

import httpx
from typing_extensions import ReadOnly

import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.http_handler import (
    AsyncHTTPHandler,
    HTTPHandler,
    _get_httpx_client,
    get_async_httpx_client,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import ImageResponse

from ..common_utils import (
    DEFAULT_MAX_POLLING_TIME,
    DEFAULT_POLLING_INTERVAL,
    BlackForestLabsError,
    assert_bfl_polling_url,
)
from .transformation import BlackForestLabsImageGenerationConfig


class _BFLTaskPayload(TypedDict, total=False):
    """The body BFL returns for a submitted or polled generation task."""

    errors: ReadOnly[object]
    polling_url: ReadOnly[str]
    status: ReadOnly[str]


class _TaskJsonResponse(Protocol):
    def json(self) -> _BFLTaskPayload: ...


def _task_payload(response: _TaskJsonResponse) -> _BFLTaskPayload:
    """The JSON body of a BFL task submission or poll response."""
    return response.json()


class BlackForestLabsImageGeneration:
    """
    Black Forest Labs Image Generation handler.

    Handles the HTTP requests and polling logic, delegating data transformation
    to the BlackForestLabsImageGenerationConfig class.
    """

    def __init__(self):
        self.config = BlackForestLabsImageGenerationConfig()

    def image_generation(
        self,
        model: str,
        prompt: str,
        model_response: ImageResponse,
        optional_params: dict,
        litellm_params: GenericLiteLLMParams | dict,
        logging_obj: LiteLLMLoggingObj,
        timeout: float | httpx.Timeout | None,
        extra_headers: Mapping[str, str] | None = None,
        client: HTTPHandler | AsyncHTTPHandler | None = None,
        aimg_generation: bool = False,
    ) -> ImageResponse | Coroutine[object, object, ImageResponse]:
        """
        Main entry point for image generation requests.

        Args:
            model: The model to use (e.g., "black_forest_labs/flux-pro-1.1")
            prompt: The text prompt for image generation
            model_response: ImageResponse object to populate
            optional_params: Optional parameters for the request
            litellm_params: LiteLLM parameters including api_key, api_base
            logging_obj: Logging object
            timeout: Request timeout
            extra_headers: Additional headers
            client: HTTP client to use
            aimg_generation: If True, return async coroutine

        Returns:
            ImageResponse or coroutine if aimg_generation=True
        """
        # Handle litellm_params as dict or object
        if isinstance(litellm_params, dict):
            api_key = litellm_params.get("api_key")
            api_base = litellm_params.get("api_base")
            litellm_params_dict = litellm_params
        else:
            api_key = litellm_params.api_key
            api_base = litellm_params.api_base
            litellm_params_dict = dict(litellm_params)

        if aimg_generation:
            return self.async_image_generation(
                model=model,
                prompt=prompt,
                model_response=model_response,
                optional_params=optional_params,
                litellm_params=litellm_params,
                logging_obj=logging_obj,
                timeout=timeout,
                extra_headers=extra_headers,
                client=client if isinstance(client, AsyncHTTPHandler) else None,
            )

        # Sync version
        if client is None or not isinstance(client, HTTPHandler):
            sync_client = _get_httpx_client()
        else:
            sync_client = client

        # Validate environment and get headers
        headers: Final = self.config.validate_environment(
            api_key=api_key,
            headers={},
            model=model,
            messages=[],
            optional_params=optional_params,
            litellm_params=litellm_params_dict,
        )
        if extra_headers:
            headers.update(extra_headers)

        # Get complete URL
        complete_url: Final = self.config.get_complete_url(
            api_base=api_base,
            api_key=api_key,
            model=model,
            optional_params=optional_params,
            litellm_params=litellm_params_dict,
        )

        # Transform request
        data: Final = self.config.transform_image_generation_request(
            model=model,
            prompt=prompt,
            optional_params=optional_params,
            litellm_params=litellm_params_dict,
            headers=headers,
        )

        # Logging
        logging_obj.pre_call(
            input=prompt,
            api_key="",
            additional_args={
                "complete_input_dict": data,
                "api_base": complete_url,
                "headers": headers,
            },
        )

        # Make initial request
        try:
            response: Final = sync_client.post(
                url=complete_url,
                headers=headers,
                json=data,
                timeout=timeout,
            )
        except Exception as e:
            raise BlackForestLabsError(
                status_code=500,
                message=f"Request failed: {e}",
            )

        # Poll for result
        final_response: Final = self._poll_for_result_sync(
            initial_response=response,
            headers=headers,
            sync_client=sync_client,
        )

        # Transform response
        return self.config.transform_image_generation_response(
            model=model,
            raw_response=final_response,
            model_response=model_response,
            logging_obj=logging_obj,
            request_data=data,
            optional_params=optional_params,
            litellm_params=litellm_params_dict,
            encoding=None,
        )

    async def async_image_generation(
        self,
        model: str,
        prompt: str,
        model_response: ImageResponse,
        optional_params: dict,
        litellm_params: GenericLiteLLMParams | dict,
        logging_obj: LiteLLMLoggingObj,
        timeout: float | httpx.Timeout | None,
        extra_headers: Mapping[str, str] | None = None,
        client: AsyncHTTPHandler | None = None,
    ) -> ImageResponse:
        """
        Async version of image generation.
        """
        # Handle litellm_params as dict or object
        if isinstance(litellm_params, dict):
            api_key = litellm_params.get("api_key")
            api_base = litellm_params.get("api_base")
            litellm_params_dict = litellm_params
        else:
            api_key = litellm_params.api_key
            api_base = litellm_params.api_base
            litellm_params_dict = dict(litellm_params)

        if client is None:
            async_client = get_async_httpx_client(
                llm_provider=litellm.LlmProviders.BLACK_FOREST_LABS,
            )
        else:
            async_client = client

        # Validate environment and get headers
        headers: Final = self.config.validate_environment(
            api_key=api_key,
            headers={},
            model=model,
            messages=[],
            optional_params=optional_params,
            litellm_params=litellm_params_dict,
        )
        if extra_headers:
            headers.update(extra_headers)

        # Get complete URL
        complete_url: Final = self.config.get_complete_url(
            api_base=api_base,
            api_key=api_key,
            model=model,
            optional_params=optional_params,
            litellm_params=litellm_params_dict,
        )

        # Transform request
        data: Final = self.config.transform_image_generation_request(
            model=model,
            prompt=prompt,
            optional_params=optional_params,
            litellm_params=litellm_params_dict,
            headers=headers,
        )

        # Logging
        logging_obj.pre_call(
            input=prompt,
            api_key="",
            additional_args={
                "complete_input_dict": data,
                "api_base": complete_url,
                "headers": headers,
            },
        )

        # Make initial request
        try:
            response: Final = await async_client.post(
                url=complete_url,
                headers=headers,
                json=data,
                timeout=timeout,
            )
        except Exception as e:
            raise BlackForestLabsError(
                status_code=500,
                message=f"Request failed: {e}",
            )

        # Poll for result
        final_response: Final = await self._poll_for_result_async(
            initial_response=response,
            headers=headers,
            async_client=async_client,
        )

        # Transform response
        return self.config.transform_image_generation_response(
            model=model,
            raw_response=final_response,
            model_response=model_response,
            logging_obj=logging_obj,
            request_data=data,
            optional_params=optional_params,
            litellm_params=litellm_params_dict,
            encoding=None,
        )

    def _poll_for_result_sync(
        self,
        initial_response: httpx.Response,
        headers: dict,
        sync_client: HTTPHandler,
        max_wait: float = DEFAULT_MAX_POLLING_TIME,
        interval: float = DEFAULT_POLLING_INTERVAL,
        timeout: float | httpx.Timeout | None = None,
    ) -> httpx.Response:
        """
        Poll BFL API until result is ready (sync version).
        """
        # Validate initial response status code
        if initial_response.status_code >= 400:
            raise BlackForestLabsError(
                status_code=initial_response.status_code,
                message=f"BFL initial request failed: {initial_response.text}",
            )

        # Parse initial response to get polling URL
        try:
            response_data: Final = _task_payload(initial_response)
        except Exception as e:
            raise BlackForestLabsError(
                status_code=initial_response.status_code,
                message=f"Error parsing initial response: {e}",
            )

        # Check for immediate errors
        if "errors" in response_data:
            raise BlackForestLabsError(
                status_code=initial_response.status_code,
                message=f"BFL error: {response_data['errors']}",
            )

        polling_url: Final = response_data.get("polling_url")
        if not polling_url:
            raise BlackForestLabsError(
                status_code=500,
                message="No polling_url in BFL response",
            )

        # Reject polling URLs that don't belong to BFL-controlled infrastructure.
        # BFL uses regional subdomains (e.g. gateway.bfl.ai) that differ from the
        # submission host (api.bfl.ai), so we validate against the registered
        # domain rather than doing a strict same-origin check. VERIA-51.
        assert_bfl_polling_url(polling_url)

        # Get just the auth header for polling
        polling_headers: Final = {"x-key": headers.get("x-key", "")}

        start_time: Final = time.time()
        verbose_logger.debug("BFL starting sync polling at %s", polling_url)

        while time.time() - start_time < max_wait:
            response = sync_client.get(
                url=polling_url,
                headers=polling_headers,
            )

            if response.status_code != 200:
                raise BlackForestLabsError(
                    status_code=response.status_code,
                    message=f"Polling failed: {response.text}",
                )

            data = _task_payload(response)
            status = data.get("status")

            verbose_logger.debug("BFL poll status: %s", status)

            if status == "Ready":
                return response
            elif status in [
                "Error",
                "Failed",
                "Content Moderated",
                "Request Moderated",
            ]:
                raise BlackForestLabsError(
                    status_code=400,
                    message=f"Image generation failed: {status}",
                )

            time.sleep(interval)

        raise BlackForestLabsError(
            status_code=408,
            message=f"Polling timed out after {max_wait} seconds",
        )

    async def _poll_for_result_async(
        self,
        initial_response: httpx.Response,
        headers: dict,
        async_client: AsyncHTTPHandler,
        max_wait: float = DEFAULT_MAX_POLLING_TIME,
        interval: float = DEFAULT_POLLING_INTERVAL,
        timeout: float | httpx.Timeout | None = None,
    ) -> httpx.Response:
        """
        Poll BFL API until result is ready (async version).
        """
        # Validate initial response status code
        if initial_response.status_code >= 400:
            raise BlackForestLabsError(
                status_code=initial_response.status_code,
                message=f"BFL initial request failed: {initial_response.text}",
            )

        # Parse initial response to get polling URL
        try:
            response_data: Final = _task_payload(initial_response)
        except Exception as e:
            raise BlackForestLabsError(
                status_code=initial_response.status_code,
                message=f"Error parsing initial response: {e}",
            )

        # Check for immediate errors
        if "errors" in response_data:
            raise BlackForestLabsError(
                status_code=initial_response.status_code,
                message=f"BFL error: {response_data['errors']}",
            )

        polling_url: Final = response_data.get("polling_url")
        if not polling_url:
            raise BlackForestLabsError(
                status_code=500,
                message="No polling_url in BFL response",
            )

        # Reject polling URLs that don't belong to BFL-controlled infrastructure.
        # BFL uses regional subdomains (e.g. gateway.bfl.ai) that differ from the
        # submission host (api.bfl.ai), so we validate against the registered
        # domain rather than doing a strict same-origin check. VERIA-51.
        assert_bfl_polling_url(polling_url)

        # Get just the auth header for polling
        polling_headers: Final = {"x-key": headers.get("x-key", "")}

        start_time: Final = time.time()
        verbose_logger.debug("BFL starting async polling at %s", polling_url)

        while time.time() - start_time < max_wait:
            response = await async_client.get(
                url=polling_url,
                headers=polling_headers,
            )

            if response.status_code != 200:
                raise BlackForestLabsError(
                    status_code=response.status_code,
                    message=f"Polling failed: {response.text}",
                )

            data = _task_payload(response)
            status = data.get("status")

            verbose_logger.debug("BFL poll status: %s", status)

            if status == "Ready":
                return response
            elif status in [
                "Error",
                "Failed",
                "Content Moderated",
                "Request Moderated",
            ]:
                raise BlackForestLabsError(
                    status_code=400,
                    message=f"Image generation failed: {status}",
                )

            await asyncio.sleep(interval)

        raise BlackForestLabsError(
            status_code=408,
            message=f"Polling timed out after {max_wait} seconds",
        )


# Singleton instance for use in images/main.py
bfl_image_generation: Final = BlackForestLabsImageGeneration()
