import asyncio
from typing import Any

import httpx
from google.genai import errors as genai_errors
from google.genai import types as genai_types
from google.genai.client import Client as GoogleGenAiClient
from typing_extensions import override

from pipelex import log
from pipelex.base_exceptions import PipelexError
from pipelex.cogt.exceptions import ImgGenGenerationError, InferenceErrorCategory, SdkTypeError
from pipelex.cogt.image.generated_image import GeneratedImageRawDetails
from pipelex.cogt.image.image_size import ImageSize
from pipelex.cogt.img_gen.img_gen_job import ImgGenJob
from pipelex.cogt.img_gen.img_gen_worker_abstract import ImgGenWorkerAbstract
from pipelex.cogt.inference.error_classification import UserAction, UserActionKind, extract_google_metadata
from pipelex.cogt.inference.error_classify import classify_inference_error
from pipelex.cogt.inference.error_render import InferenceErrorFamily, render_inference_error
from pipelex.cogt.model_backends.model_spec import InferenceModelSpec
from pipelex.plugins.google.google_factory import GoogleFactory
from pipelex.plugins.google.google_img_gen_factory import GoogleImgGenFactory
from pipelex.reporting.reporting_protocol import ReportingProtocol


class GoogleImgGenWorkerError(PipelexError):
    """Base exception for Google Image Generation Worker errors."""


class GoogleImgGenWorker(ImgGenWorkerAbstract):
    """Worker for generating images using Google Gemini Image models (Nano Banana)."""

    def __init__(
        self,
        sdk_instance: Any,
        inference_model: InferenceModelSpec,
        reporting_delegate: ReportingProtocol | None = None,
    ):
        super().__init__(inference_model=inference_model, reporting_delegate=reporting_delegate)

        if not isinstance(sdk_instance, GoogleGenAiClient):
            msg = f"Provided ImgGen sdk_instance is not of type GoogleGenAiClient: it's a '{type(sdk_instance)}'"
            raise SdkTypeError(msg)

        self.genai_client: GoogleGenAiClient = sdk_instance
        self.genai_async_client = sdk_instance.aio

        # Capture the event loop at creation time if one is running
        self._event_loop: asyncio.AbstractEventLoop | None
        try:
            self._event_loop = asyncio.get_running_loop()
        except RuntimeError:
            # No running loop at creation time
            self._event_loop = None

    @override
    def teardown(self):
        """Close the async client to free resources."""
        try:
            # First, try to use the loop captured at creation time if it's still running
            if self._event_loop is not None and self._event_loop.is_running():
                # Schedule cleanup on the captured loop and store reference to prevent garbage collection
                task = self._event_loop.create_task(self.genai_async_client.aclose())
                # Add a callback to log any errors that occur during cleanup
                task.add_done_callback(lambda t: log.debug(f"Google async client cleanup error: {t.exception()}") if t.exception() else None)
                log.verbose("Scheduled Google async client cleanup on captured event loop")
                return

            # Otherwise, try to get the current running loop
            try:
                current_loop = asyncio.get_running_loop()
                # Schedule cleanup on the current running loop and store reference to prevent garbage collection
                task = current_loop.create_task(self.genai_async_client.aclose())
                # Add a callback to log any errors that occur during cleanup
                task.add_done_callback(lambda t: log.debug(f"Google async client cleanup error: {t.exception()}") if t.exception() else None)
                log.verbose("Scheduled Google async client cleanup on current event loop")
            except RuntimeError:
                # No running event loop, we can safely use asyncio.run()
                try:
                    asyncio.run(self.genai_async_client.aclose())
                    log.verbose("Closed Google async client using asyncio.run()")
                except Exception as exc:  # noqa: BLE001
                    # Best-effort: asyncio.run() runs aclose(), whose failure surface is not enumerable; teardown must never fail.
                    log.verbose(f"Error closing Google async client during teardown: {exc}")
        except Exception as exc:  # noqa: BLE001
            # Best-effort cleanup boundary: teardown must never fail, whatever client/event-loop close throws.
            log.debug(f"Error during Google async client teardown: {exc}")

    @override
    async def _gen_image(
        self,
        img_gen_job: ImgGenJob,
    ) -> GeneratedImageRawDetails:
        """Generate a single image using Google Gemini Image API."""
        prompt_text = img_gen_job.img_gen_prompt.positive_text
        aspect_ratio_str = GoogleImgGenFactory.aspect_ratio_literal(img_gen_job.job_params.aspect_ratio)
        width, height = GoogleImgGenFactory.dimensions_for_aspect_ratio_and_size(
            model=self.inference_model.name,
            aspect_ratio=img_gen_job.job_params.aspect_ratio,
            size="1K",
        )

        # Build image config for aspect ratio
        image_config = genai_types.ImageConfig(
            aspect_ratio=aspect_ratio_str,
        )

        # Build generation config with image output
        generation_config = genai_types.GenerateContentConfig(
            response_modalities=["Image"],
            image_config=image_config,
            # seed=img_gen_job.job_params.seed,
        )

        # Generate content using async client
        try:
            response = await self.genai_async_client.models.generate_content(
                model=self.inference_model.model_id,
                contents=prompt_text,
                config=generation_config,
            )
        except (genai_errors.ServerError, genai_errors.ClientError, httpx.TransportError) as exc:
            metadata = extract_google_metadata(exc)
            classification = classify_inference_error(metadata)
            raise render_inference_error(
                metadata=metadata,
                classification=classification,
                family=InferenceErrorFamily.IMG_GEN,
                model_desc=self.inference_model.desc,
                model_handle=self.inference_model.name,
            ) from exc

        usage_metadata: genai_types.GenerateContentResponseUsageMetadata | None = response.usage_metadata
        if not usage_metadata:
            log.warning("No usage metadata returned from Google")

        if usage_metadata and (img_gen_tokens_usage := img_gen_job.job_report.img_gen_tokens_usage):
            request_tokens = GoogleFactory.extract_token_usage(usage_metadata)
            for token_category, nb_tokens in request_tokens.items():
                img_gen_tokens_usage.nb_tokens_by_category[token_category] = (
                    img_gen_tokens_usage.nb_tokens_by_category.get(token_category, 0) + nb_tokens
                )

        # Extract image from response
        if not response.candidates:
            msg = f"No candidates returned from model: {self.inference_model.desc}"
            raise ImgGenGenerationError(
                msg,
                error_category=InferenceErrorCategory.CONTENT,
                user_action=UserAction(
                    kind=UserActionKind.CHANGE_INPUT,
                    detail="Google returned no image candidates — try rephrasing the prompt or using a different model",
                ),
                provider_metadata=None,
            )

        candidate = response.candidates[0]
        if not candidate.content or not candidate.content.parts:
            msg = f"No content parts in response from model: {self.inference_model.desc}"
            raise ImgGenGenerationError(
                msg,
                error_category=InferenceErrorCategory.CONTENT,
                user_action=UserAction(
                    kind=UserActionKind.CHANGE_INPUT,
                    detail="Google returned an empty response — try rephrasing the prompt or using a different model",
                ),
                provider_metadata=None,
            )

        # Look for image data in response parts
        for part in candidate.content.parts:
            if part.inline_data is not None:
                image_bytes = part.inline_data.data
                if image_bytes is None:
                    continue
                mime_type = part.inline_data.mime_type
                if not mime_type:
                    msg = "No mime type returned from Google"
                    raise ImgGenGenerationError(
                        msg,
                        error_category=InferenceErrorCategory.UNKNOWN,
                        user_action=UserAction(
                            kind=UserActionKind.CHANGE_MODEL,
                            detail="Google returned image data without a mime type — try a different model",
                        ),
                        provider_metadata=None,
                    )
                return GeneratedImageRawDetails(
                    actual_bytes=image_bytes,
                    size=ImageSize(width=width, height=height),
                    mime_type=mime_type,
                )

        msg = f"No image data in response from model: {self.inference_model.desc}"
        raise ImgGenGenerationError(
            msg,
            error_category=InferenceErrorCategory.CONTENT,
            user_action=UserAction(
                kind=UserActionKind.CHANGE_INPUT,
                detail="Google returned no image data — try rephrasing the prompt or using a different model",
            ),
            provider_metadata=None,
        )

    @override
    async def _gen_image_list(
        self,
        img_gen_job: ImgGenJob,
        nb_images: int,
    ) -> list[GeneratedImageRawDetails]:
        """Generate multiple images by calling _gen_image multiple times.

        Google Gemini Image API does not support generating multiple images in a single call,
        so we generate them sequentially.
        """
        generated_images: list[GeneratedImageRawDetails] = []
        # TODO: async gen images in parallel
        for _ in range(nb_images):
            image = await self._gen_image(img_gen_job=img_gen_job)
            generated_images.append(image)
        return generated_images
