"""FAL.ai image generation backend.

Wraps the FAL catalog (FLUX 2, Z-Image, Nano Banana, GPT Image 1.5, Recraft,
Imagen 4, Qwen, Ideogram, …) as an :class:`ImageGenProvider`.

The heavy lifting — model catalog, payload construction, request submission,
managed-Nous-gateway selection, Clarity Upscaler chaining — lives in
:mod:`tools.image_generation_tool`. This plugin reaches into that module via
call-time indirection (``import tools.image_generation_tool as _it``) so the
existing tests keep patching ``image_tool.*`` unchanged, and there is exactly
one canonical FAL code path on disk — the plugin is a registration adapter.
"""

from __future__ import annotations

import json
import logging
from typing import Any, Dict, List, Optional

from agent.image_gen_provider import (
    DEFAULT_ASPECT_RATIO,
    ImageGenProvider,
    resolve_aspect_ratio,
)
from plugins.image_gen._common import api_key_setup_schema, catalog_rows

logger = logging.getLogger(__name__)

_PASSTHROUGH_KWARGS = (
    "num_inference_steps", "guidance_scale", "num_images", "output_format", "seed", "upscale",
)


class FalImageGenProvider(ImageGenProvider):
    """FAL.ai backend delegating to ``tools.image_generation_tool`` at call time."""

    @property
    def name(self) -> str:
        return "fal"

    @property
    def display_name(self) -> str:
        return "FAL.ai"

    def is_available(self) -> bool:
        # Direct FAL_KEY or a managed Nous fal-queue origin; both checks live in
        # the legacy module so this provider tracks whatever logic ships there.
        import tools.image_generation_tool as _it
        try:
            return bool(_it.check_fal_api_key())
        except Exception:  # noqa: BLE001 — never break the picker
            return False

    def list_models(self) -> List[Dict[str, Any]]:
        import tools.image_generation_tool as _it
        return catalog_rows(_it.FAL_MODELS)

    def default_model(self) -> Optional[str]:
        import tools.image_generation_tool as _it
        return _it.DEFAULT_MODEL

    def get_setup_schema(self) -> Dict[str, Any]:
        return api_key_setup_schema(
            "FAL.ai", "paid",
            "Pick from flux-2-klein, flux-2-pro, gpt-image, nano-banana-2, nano-banana-pro, etc. — text-to-image & image editing",
            key="FAL_KEY", prompt="FAL API key", url="https://fal.ai/dashboard/keys",
        )

    def capabilities(self) -> Dict[str, Any]:
        # Image-to-image depends on the currently selected FAL model (each entry
        # declares an edit_endpoint or not); Clarity Upscaler chains on request
        # for any model.
        import tools.image_generation_tool as _it

        try:
            _model_id, meta = _it._resolve_fal_model()
        except Exception:  # noqa: BLE001
            return {"modalities": ["text"], "max_reference_images": 0}
        if meta.get("edit_endpoint"):
            return {
                "modalities": ["text", "image"],
                "max_reference_images": int(meta.get("max_reference_images") or 1),
                "supports_upscale": True,
            }
        return {"modalities": ["text"], "max_reference_images": 0, "supports_upscale": True}

    def generate(
        self,
        prompt: str,
        aspect_ratio: str = DEFAULT_ASPECT_RATIO,
        *,
        image_url: Optional[str] = None,
        reference_image_urls: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> Dict[str, Any]:
        """Forward to :func:`tools.image_generation_tool.image_generate_tool` and
        reshape its JSON-string response into the provider-ABC dict."""
        import tools.image_generation_tool as _it

        aspect = resolve_aspect_ratio(aspect_ratio)
        passthrough = {
            key: kwargs[key] for key in _PASSTHROUGH_KWARGS
            if key in kwargs and kwargs[key] is not None
        }
        # Only forward image-to-image inputs when supplied, so a plain
        # text-to-image call delegates exactly as before (no noisy None kwargs).
        if image_url is not None:
            passthrough["image_url"] = image_url
        if reference_image_urls is not None:
            passthrough["reference_image_urls"] = reference_image_urls

        try:
            raw = _it.image_generate_tool(prompt=prompt, aspect_ratio=aspect, **passthrough)
        except Exception as exc:  # noqa: BLE001 — never raise out of generate
            logger.warning("FAL image_generate_tool raised: %s", exc, exc_info=True)
            return {
                "success": False,
                "image": None,
                "error": f"FAL image generation failed: {exc}",
                "error_type": type(exc).__name__,
                "provider": "fal",
                "prompt": prompt,
                "aspect_ratio": aspect,
            }

        try:
            response = json.loads(raw) if isinstance(raw, str) else raw
        except Exception:  # noqa: BLE001
            response = {"success": False, "image": None, "error": "Invalid JSON from FAL pipeline"}

        if not isinstance(response, dict):
            response = {
                "success": False,
                "image": None,
                "error": "FAL pipeline returned a non-dict response",
                "error_type": "provider_contract",
            }

        # Stamp the uniform shape declared in ``agent.image_gen_provider``; the
        # legacy pipeline resolves the model internally, so query it after the fact.
        response.setdefault("provider", "fal")
        response.setdefault("prompt", prompt)
        response.setdefault("aspect_ratio", aspect)
        if "model" not in response:
            try:
                response["model"] = _it._resolve_fal_model()[0]
            except Exception:  # noqa: BLE001
                pass
        return response


def register(ctx) -> None:
    """Plugin entry point — wire ``FalImageGenProvider`` into the registry."""
    ctx.register_image_gen_provider(FalImageGenProvider())