from __future__ import annotations

import os
import io
from typing import Optional, Dict, Any
from PIL import Image

from .image_generation import (
    ImageGenerationOptions,
    ImageGenerationResult,
)
from utils.logger_utils import get_service_logger

try:
    from huggingface_hub import InferenceClient
    HF_HUB_AVAILABLE = True
except ImportError:
    HF_HUB_AVAILABLE = False


logger = get_service_logger("image_editing.facade")


DEFAULT_IMAGE_EDIT_MODEL = os.getenv(
    "HF_IMAGE_EDIT_MODEL",
    "Qwen/Qwen-Image-Edit",
)


def _select_provider(explicit: Optional[str]) -> str:
    """Select provider for image editing. Defaults to huggingface with fal-ai."""
    if explicit:
        return explicit
    # Default to huggingface for image editing (best support for image-to-image)
    return "huggingface"


def _get_provider_client(provider_name: str, api_key: Optional[str] = None):
    """Get InferenceClient for the specified provider."""
    if not HF_HUB_AVAILABLE:
        raise RuntimeError("huggingface_hub is not installed. Install with: pip install huggingface_hub")
    
    if provider_name == "huggingface":
        api_key = api_key or os.getenv("HF_TOKEN")
        if not api_key:
            raise RuntimeError("HF_TOKEN is required for Hugging Face image editing")
        # Use fal-ai provider for fast inference
        return InferenceClient(provider="fal-ai", api_key=api_key)
    
    raise ValueError(f"Unknown image editing provider: {provider_name}")


def edit_image(
    input_image_bytes: bytes,
    prompt: str,
    options: Optional[Dict[str, Any]] = None,
    user_id: Optional[str] = None,
    mask_bytes: Optional[bytes] = None,
) -> ImageGenerationResult:
    """Edit image with pre-flight validation.
    
    Args:
        input_image_bytes: Input image as bytes (PNG/JPEG)
        prompt: Natural language prompt describing desired edits (e.g., "Turn the cat into a tiger")
        options: Image editing options (provider, model, etc.)
        user_id: User ID for subscription checking (optional, but required for validation)
        mask_bytes: Optional mask image bytes for selective editing (grayscale, white=edit, black=preserve)
    
    Returns:
        ImageGenerationResult with edited image bytes and metadata
        
    Best Practices for Prompts:
    - Use clear, specific language describing desired changes
    - Describe what should change and what should remain
    - Examples: "Turn the cat into a tiger", "Change background to forest", 
                "Make it look like a watercolor painting"
    
    Note: Mask support depends on the specific model. Some models may ignore the mask parameter.
    """
    # PRE-FLIGHT VALIDATION: Validate image editing before API call
    # MUST happen BEFORE any API calls - return immediately if validation fails
    if user_id:
        from services.database import get_db
        from services.subscription import PricingService
        from services.subscription.preflight_validator import validate_image_editing_operations
        from fastapi import HTTPException
        
        logger.info(f"[Image Editing] 🔍 Starting pre-flight validation for user_id={user_id}")
        db = next(get_db())
        try:
            pricing_service = PricingService(db)
            # Raises HTTPException immediately if validation fails - frontend gets immediate response
            validate_image_editing_operations(
                pricing_service=pricing_service,
                user_id=user_id
            )
            logger.info(f"[Image Editing] ✅ Pre-flight validation passed for user_id={user_id} - proceeding with image editing")
        except HTTPException as http_ex:
            # Re-raise immediately - don't proceed with API call
            logger.error(f"[Image Editing] ❌ Pre-flight validation failed for user_id={user_id} - blocking API call: {http_ex.detail}")
            raise
        finally:
            db.close()
    else:
        logger.warning(f"[Image Editing] ⚠️ No user_id provided - skipping pre-flight validation (this should not happen in production)")
    
    # Validate input
    if not input_image_bytes:
        raise ValueError("input_image_bytes is required")
    if not prompt or not prompt.strip():
        raise ValueError("prompt is required for image editing")
    
    opts = options or {}
    provider_name = _select_provider(opts.get("provider"))
    model = opts.get("model") or DEFAULT_IMAGE_EDIT_MODEL
    
    logger.info(f"[Image Editing] Editing image via provider={provider_name} model={model}")
    
    # Get provider client
    client = _get_provider_client(provider_name, opts.get("api_key"))
    
    # Prepare parameters for image-to-image
    params: Dict[str, Any] = {}
    if opts.get("guidance_scale") is not None:
        params["guidance_scale"] = opts.get("guidance_scale")
    if opts.get("steps") is not None:
        params["num_inference_steps"] = opts.get("steps")
    if opts.get("seed") is not None:
        params["seed"] = opts.get("seed")
    
    try:
        # Convert input image bytes to PIL Image for validation
        input_image = Image.open(io.BytesIO(input_image_bytes))
        width = input_image.width
        height = input_image.height
        
        # Convert mask bytes to PIL Image if provided
        mask_image = None
        if mask_bytes:
            try:
                mask_image = Image.open(io.BytesIO(mask_bytes)).convert("L")  # Convert to grayscale
                # Ensure mask dimensions match input image
                if mask_image.size != input_image.size:
                    logger.warning(f"[Image Editing] Mask size {mask_image.size} doesn't match image size {input_image.size}, resizing mask")
                    mask_image = mask_image.resize(input_image.size, Image.Resampling.LANCZOS)
            except Exception as e:
                logger.warning(f"[Image Editing] Failed to process mask image: {e}, continuing without mask")
                mask_image = None
        
        # Use image_to_image method from Hugging Face InferenceClient
        # This follows the pattern from the Hugging Face documentation
        # Docs: https://huggingface.co/docs/inference-providers/en/guides/image-editor
        # Note: Mask support depends on the model - some models may ignore it
        call_params = params.copy()
        if mask_image:
            call_params["mask_image"] = mask_image
            logger.info("[Image Editing] Using mask for selective editing")
        
        edited_image: Image.Image = client.image_to_image(
            image=input_image,
            prompt=prompt.strip(),
            model=model,
            **call_params,
        )
        
        # Convert edited image back to bytes
        with io.BytesIO() as buf:
            edited_image.save(buf, format="PNG")
            edited_image_bytes = buf.getvalue()
        
        logger.info(f"[Image Editing] ✅ Successfully edited image: {len(edited_image_bytes)} bytes")
        
        return ImageGenerationResult(
            image_bytes=edited_image_bytes,
            width=edited_image.width,
            height=edited_image.height,
            provider="huggingface",
            model=model,
            seed=opts.get("seed"),
            metadata={
                "provider": "fal-ai",
                "operation": "image_editing",
                "original_width": width,
                "original_height": height,
            },
        )
    except Exception as e:
        logger.error(f"[Image Editing] ❌ Error editing image: {e}", exc_info=True)
        raise RuntimeError(f"Image editing failed: {str(e)}")

