"""Image generation tool for diagrams, overlays, and illustrations.

.. deprecated::
    Use ``image_selector`` instead. This monolithic tool has been replaced by
    the selector/provider pattern: ``image_selector`` routes to per-provider
    tools (flux_image, openai_image, recraft_image, local_diffusion,
    pexels_image, pixabay_image). This file is kept for backwards
    compatibility and will be removed in a future release.

Supports cloud API providers (FLUX via fal.ai/Replicate, OpenAI DALL-E)
and local Stable Diffusion via diffusers. Reports unavailable with
install instructions when no provider is configured.
"""

from __future__ import annotations

import json
import os
import time
from pathlib import Path
from typing import Any, Optional

from tools.base_tool import (
    BaseTool,
    Determinism,
    ExecutionMode,
    ResourceProfile,
    RetryPolicy,
    ToolResult,
    ToolRuntime,
    ToolStability,
    ToolStatus,
    ToolTier,
)


class ImageGen(BaseTool):
    name = "image_gen"
    version = "0.1.0"
    tier = ToolTier.CORE
    capability = "image_generation"
    provider = "multi"
    stability = ToolStability.EXPERIMENTAL
    execution_mode = ExecutionMode.SYNC
    determinism = Determinism.SEEDED
    runtime = ToolRuntime.HYBRID  # API (DALL-E/FLUX) or local (diffusers)

    dependencies = []  # checked dynamically based on provider
    install_instructions = (
        "Set one of these environment variables:\n"
        "  OPENAI_API_KEY — for DALL-E 3\n"
        "  FAL_KEY — for FLUX via fal.ai\n"
        "Or install diffusers for local generation:\n"
        "  pip install diffusers transformers accelerate torch"
    )
    agent_skills = ["flux-best-practices", "bfl-api"]

    capabilities = [
        "generate_image",
        "generate_diagram_overlay",
        "generate_illustration",
    ]
    best_for = [
        "DEPRECATED — prefer image_selector which routes to per-provider tools "
        "(flux_image, openai_image, recraft_image, grok_image, local_diffusion, "
        "pexels_image, pixabay_image).",
        "Kept only for backwards compatibility. New code should not call this.",
    ]
    not_good_for = [
        "New production code — use image_selector instead.",
        "Picking a specific provider (use the per-provider tool directly).",
    ]

    input_schema = {
        "type": "object",
        "required": ["prompt"],
        "properties": {
            "prompt": {"type": "string"},
            "negative_prompt": {"type": "string", "default": ""},
            "width": {"type": "integer", "default": 1024},
            "height": {"type": "integer", "default": 1024},
            "provider": {
                "type": "string",
                "enum": ["openai", "flux", "local"],
                "description": "Auto-detected if not specified",
            },
            "model": {"type": "string"},
            "seed": {"type": "integer"},
            "output_path": {"type": "string"},
        },
    }

    resource_profile = ResourceProfile(
        cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=100, network_required=True
    )
    retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
    idempotency_key_fields = ["prompt", "width", "height", "seed"]
    side_effects = ["writes image file to output_path", "calls external API"]
    user_visible_verification = [
        "Inspect generated image for relevance and quality",
    ]

    def get_status(self) -> ToolStatus:
        provider = self._detect_provider()
        if provider:
            return ToolStatus.AVAILABLE
        return ToolStatus.UNAVAILABLE

    def _detect_provider(self) -> Optional[str]:
        if os.environ.get("OPENAI_API_KEY"):
            return "openai"
        if os.environ.get("FAL_KEY") or os.environ.get("FAL_AI_API_KEY"):
            return "flux"
        try:
            import diffusers  # noqa: F401
            return "local"
        except ImportError:
            pass
        return None

    def estimate_cost(self, inputs: dict[str, Any]) -> float:
        provider = inputs.get("provider") or self._detect_provider()
        if provider == "openai":
            return 0.04  # DALL-E 3 standard
        if provider == "flux":
            return 0.03
        return 0.0  # local

    def execute(self, inputs: dict[str, Any]) -> ToolResult:
        provider = inputs.get("provider") or self._detect_provider()
        if not provider:
            return ToolResult(
                success=False,
                error="No image generation provider available. " + self.install_instructions,
            )

        start = time.time()

        try:
            if provider == "openai":
                result = self._generate_openai(inputs)
            elif provider == "flux":
                result = self._generate_flux(inputs)
            elif provider == "local":
                result = self._generate_local(inputs)
            else:
                return ToolResult(success=False, error=f"Unknown provider: {provider}")
        except Exception as e:
            return ToolResult(success=False, error=f"Generation failed: {e}")

        result.duration_seconds = round(time.time() - start, 2)
        result.cost_usd = self.estimate_cost(inputs)
        return result

    def _generate_openai(self, inputs: dict[str, Any]) -> ToolResult:
        from openai import OpenAI
        import base64

        client = OpenAI()
        prompt = inputs["prompt"]
        size = f"{inputs.get('width', 1024)}x{inputs.get('height', 1024)}"
        model = inputs.get("model", "dall-e-3")

        response = client.images.generate(
            model=model,
            prompt=prompt,
            size=size,
            n=1,
            response_format="b64_json",
        )

        image_data = base64.b64decode(response.data[0].b64_json)
        output_path = Path(inputs.get("output_path", "generated_image.png"))
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_bytes(image_data)

        return ToolResult(
            success=True,
            data={
                "provider": "openai",
                "model": model,
                "prompt": prompt,
                "output": str(output_path),
            },
            artifacts=[str(output_path)],
            model=model,
        )

    def _generate_flux(self, inputs: dict[str, Any]) -> ToolResult:
        import requests

        api_key = os.environ.get("FAL_KEY") or os.environ["FAL_AI_API_KEY"]
        prompt = inputs["prompt"]
        width = inputs.get("width", 1024)
        height = inputs.get("height", 1024)
        seed = inputs.get("seed")

        payload = {
            "prompt": prompt,
            "image_size": {"width": width, "height": height},
        }
        if seed is not None:
            payload["seed"] = seed

        response = requests.post(
            "https://fal.run/fal-ai/flux/dev",
            headers={"Authorization": f"Key {api_key}", "Content-Type": "application/json"},
            json=payload,
            timeout=120,
        )
        response.raise_for_status()
        data = response.json()

        image_url = data["images"][0]["url"]
        image_response = requests.get(image_url, timeout=60)
        image_response.raise_for_status()

        output_path = Path(inputs.get("output_path", "generated_image.png"))
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_bytes(image_response.content)

        return ToolResult(
            success=True,
            data={
                "provider": "flux",
                "prompt": prompt,
                "output": str(output_path),
                "seed": data.get("seed"),
            },
            artifacts=[str(output_path)],
            seed=data.get("seed"),
            model="flux-dev",
        )

    def _generate_local(self, inputs: dict[str, Any]) -> ToolResult:
        import torch
        from diffusers import StableDiffusionPipeline

        prompt = inputs["prompt"]
        negative = inputs.get("negative_prompt", "")
        width = inputs.get("width", 512)
        height = inputs.get("height", 512)
        seed = inputs.get("seed")
        model_id = inputs.get("model", "stabilityai/stable-diffusion-2-1-base")

        device = "cuda" if torch.cuda.is_available() else "cpu"
        dtype = torch.float16 if device == "cuda" else torch.float32

        pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype)
        pipe = pipe.to(device)

        generator = None
        if seed is not None:
            generator = torch.Generator(device=device).manual_seed(seed)

        image = pipe(
            prompt,
            negative_prompt=negative,
            width=width,
            height=height,
            generator=generator,
        ).images[0]

        output_path = Path(inputs.get("output_path", "generated_image.png"))
        output_path.parent.mkdir(parents=True, exist_ok=True)
        image.save(str(output_path))

        return ToolResult(
            success=True,
            data={
                "provider": "local",
                "model": model_id,
                "prompt": prompt,
                "output": str(output_path),
            },
            artifacts=[str(output_path)],
            seed=seed,
            model=model_id,
        )
