"""Image generation tool using fal.ai APIs."""
import os
import json
import time
import tempfile
from io import BytesIO
from typing import Optional, Literal, List

import fal_client
import requests
from pydantic import BaseModel, Field
from langchain_core.tools import tool
from PIL import Image

# Configure fal.ai
fal_client.api_key = os.getenv("FAL_KEY")

# Model mapping for pro mode
PRO_MODELS = {
    "seedream4": "fal-ai/bytedance/seedream/v4/text-to-image",
    "imagen4_ultra": "fal-ai/imagen4/preview/ultra",
    "flux-krea": "fal-ai/flux/krea",
}


class ImageGenerationInput(BaseModel):
    """Input schema for image generation."""
    prompt: str = Field(description="Detailed text description of the image to generate")
    mode: Literal["fast", "pro"] = Field(
        default="fast",
        description="Generation mode: 'fast' uses z-image turbo for quick results, 'pro' uses premium fal.ai models"
    )
    model: str = Field(
        default="flux-krea",
        description="Model for pro mode: 'seedream4', 'imagen4_ultra', 'flux-krea'. Default: flux-krea"
    )
    width: int = Field(default=720, description="Image width in pixels (256-2048)")
    height: int = Field(default=720, description="Image height in pixels (256-2048)")
    num_images: int = Field(default=4, description="Number of images to generate (1-4)")
    negative_prompt: Optional[str] = Field(
        default=None,
        description="Things to avoid in the generated image"
    )


def _download_and_save_images(image_results: list) -> List[str]:
    """Download images from URLs and save to temp files."""
    saved_paths = []
    for img_data in image_results:
        try:
            img_url = img_data.get("url") if isinstance(img_data, dict) else str(img_data)
            if not img_url:
                continue
            response = requests.get(img_url, timeout=30)
            response.raise_for_status()
            img = Image.open(BytesIO(response.content))
            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
            img.save(temp_file.name)
            saved_paths.append(temp_file.name)
        except Exception as e:
            print(f"DEBUG - Error downloading image: {e}")
    return saved_paths


def _generate_fast(prompt: str, width: int, height: int, num_images: int) -> List[str]:
    """Generate images using z-image turbo (fast mode)."""
    result = fal_client.subscribe(
        "fal-ai/z-image/turbo",
        arguments={
            "prompt": prompt,
            "image_size": {"width": width, "height": height},
            "num_images": num_images,
            "acceleration": "high",
            "num_inference_steps": 8,
        },
        with_logs=True,
    )

    if "images" in result and result["images"]:
        return _download_and_save_images(result["images"])

    raise ValueError("No images generated by z-image turbo")


def _generate_pro(prompt: str, model_id: str, width: int, height: int, num_images: int) -> List[str]:
    """Generate images using premium fal.ai models (pro mode)."""
    args = {
        "prompt": prompt,
        "image_size": {"width": width, "height": height},
        "num_images": num_images,
    }

    handler = fal_client.submit(model_id, arguments=args)
    result = handler.get()

    if "images" in result and result["images"]:
        return _download_and_save_images(result["images"])

    raise ValueError(f"No images generated by {model_id}")


@tool(args_schema=ImageGenerationInput)
def generate_images(
    prompt: str,
    mode: Literal["fast", "pro"] = "fast",
    model: str = "flux-krea",
    width: int = 720,
    height: int = 720,
    num_images: int = 4,
    negative_prompt: Optional[str] = None
) -> str:
    """Generate images from a text prompt using AI models.

    Use this tool when the user wants to create new images from scratch based on a description.
    Returns file paths to the generated images.

    Args:
        prompt: Detailed text description of the image to generate
        mode: 'fast' for quick generation (z-image turbo), 'pro' for higher quality (premium models)
        model: Model for pro mode - 'seedream4', 'imagen4_ultra', 'flux-krea'
        width: Image width in pixels
        height: Image height in pixels
        num_images: Number of images to generate (1-4)
        negative_prompt: Things to avoid in the generated image

    Returns:
        String containing the file paths of generated images
    """
    # Clamp values
    width = max(256, min(2048, width))
    height = max(256, min(2048, height))
    num_images = max(1, min(4, num_images))

    # Get model ID for pro mode
    model_id = PRO_MODELS.get(model, PRO_MODELS["flux-krea"])

    # Debug logging
    debug_info = {
        "prompt": prompt[:100] + "..." if len(prompt) > 100 else prompt,
        "mode": mode,
        "model": model if mode == "pro" else "z-image-turbo",
        "model_id": model_id if mode == "pro" else "fal-ai/z-image/turbo",
        "width": width,
        "height": height,
        "num_images": num_images,
    }
    print(f"\n{'='*80}")
    print(f"DEBUG - Image Generation Tool Called:")
    print(json.dumps(debug_info, indent=2))
    print(f"{'='*80}\n")

    st = time.time()
    generated_paths = []

    try:
        if mode == "fast":
            generated_paths = _generate_fast(prompt, width, height, num_images)
        else:
            generated_paths = _generate_pro(prompt, model_id, width, height, num_images)
    except Exception as e:
        et = time.time()
        print(f"DEBUG - Image generation FAILED after {et - st:.2f}s: {e}")
        return f"Failed to generate images: {str(e)}"

    et = time.time()

    if not generated_paths:
        print(f"DEBUG - Image generation FAILED after {et - st:.2f}s: No images returned")
        return "Failed to generate images: No images returned from API"

    print(f"DEBUG - Generated {len(generated_paths)} images in {et - st:.2f}s")
    for path in generated_paths:
        print(f"DEBUG - Image path: {path}")

    result = f"Successfully generated {len(generated_paths)} image(s): {', '.join(generated_paths)}"
    print(f"DEBUG - Image generation result:\n{result}")
    return result