import io
import os
import time
import tempfile
import hashlib
import asyncio
import concurrent.futures
from typing import Tuple, Dict
from PIL import Image
import torch
import numpy as np
import fal_client as fal
from fal_client import client
import aiohttp


class NxdifyNode:
    """
    ComfyUI node for Nxdify image generation using FAL AI Seedream 4.5.
    Takes 4 reference images (Face, Body, Breasts, Dynamic Pose) and generates variations.
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "face_image": ("IMAGE",),
                "body_image": ("IMAGE",),
                "breasts_image": ("IMAGE",),
                "dynamic_pose_image": ("IMAGE",),
                "prompt": ("STRING", {
                    "multiline": True,
                    "default": ""
                }),
                "fal_api_key": ("STRING", {
                    "default": "",
                    "password": True
                }),
                "quality": (["auto_4K", "auto_2K"], {
                    "default": "auto_4K"
                }),
            }
        }
    
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "execute"
    CATEGORY = "image/generation"
    
    MAX_IMAGE_SIZE = 5 * 1024 * 1024  # 5MB
    MAX_CONCURRENT = 8
    
    # Class-level cache for uploaded reference image URLs (hash -> URL)
    _image_url_cache: Dict[str, str] = {}
    
    def compress_image_bytes_max(self, image_bytes: bytes, max_bytes: int) -> bytes:
        """
        Compress image to fit under max_bytes.
        Strategy:
        1. Try reducing JPEG quality (start at 92, down to 52)
        2. If still too large, downscale image (start at 100%, down to 45%)
        3. Repeat until under limit or minimums reached
        """
        if len(image_bytes) <= max_bytes:
            return image_bytes
        
        # Convert to PIL Image
        img = Image.open(io.BytesIO(image_bytes))
        img = img.convert("RGB")
        base_w, base_h = img.size
        
        quality = 92
        scale = 1.0
        
        for _ in range(20):  # Max 20 iterations
            w = max(1, int(base_w * scale))
            h = max(1, int(base_h * scale))
            
            # Resize if needed
            working = img if (w == base_w and h == base_h) else img.resize((w, h), Image.Resampling.LANCZOS)
            
            # Save as JPEG
            buf = io.BytesIO()
            working.save(buf, format="JPEG", quality=quality, optimize=True)
            data = buf.getvalue()
            
            if len(data) <= max_bytes:
                return data
            
            # Reduce quality first
            if quality > 52:
                quality = max(52, quality - 10)
                continue
            
            # Then downscale
            if scale > 0.45:
                scale = scale * 0.85
                quality = 92  # Reset quality
                continue
            
            # Can't compress further
            return data
        
        return image_bytes
    
    def tensor_to_bytes(self, tensor: torch.Tensor) -> bytes:
        """Convert ComfyUI image tensor to JPEG bytes."""
        # ComfyUI IMAGE tensors are in BHWC format (batch, height, width, channels)
        # Remove batch dimension to get HWC
        if len(tensor.shape) == 4:
            img_array = tensor[0].cpu().numpy()  # Shape: (height, width, channels)
        else:
            img_array = tensor.cpu().numpy()
        
        # Ensure values are in 0-255 range and convert to uint8
        img_array = (np.clip(img_array, 0.0, 1.0) * 255.0).astype(np.uint8)
        
        # Handle alpha channel if present
        if img_array.shape[2] == 4:
            # Convert RGBA to RGB with white background
            alpha = img_array[:, :, 3:4].astype(np.float32) / 255.0
            rgb = img_array[:, :, :3].astype(np.float32)
            img_array = (rgb * alpha + 255 * (1 - alpha)).astype(np.uint8)
        elif img_array.shape[2] == 1:
            # Handle grayscale - convert to RGB
            img_array = np.repeat(img_array, 3, axis=2)
        
        # Convert to PIL Image
        img = Image.fromarray(img_array)
        
        # Convert to RGB if needed
        if img.mode != "RGB":
            img = img.convert("RGB")
        
        # Save to bytes
        buf = io.BytesIO()
        img.save(buf, format="JPEG", quality=95, optimize=True)
        return buf.getvalue()
    
    def _compute_image_hash(self, image_bytes: bytes) -> str:
        """Compute SHA256 hash of image bytes for caching."""
        return hashlib.sha256(image_bytes).hexdigest()
    
    def _upload_file_sync(self, tmp_path: str) -> str:
        """Synchronous wrapper for upload_file to use with asyncio.to_thread."""
        return fal.upload_file(tmp_path)
    
    async def upload_ref_with_retry(self, image_bytes: bytes, use_cache: bool = True, max_attempts: int = 3) -> str:
        """Upload image with retry on timeout. Optionally use cache to avoid re-uploading."""
        upload_start = time.time()
        original_size = len(image_bytes)
        
        # Check cache first if enabled
        if use_cache:
            image_hash = self._compute_image_hash(image_bytes)
            if image_hash in self._image_url_cache:
                print(f"[Nxdify] Image found in cache (hash: {image_hash[:16]}...), skipping upload")
                return self._image_url_cache[image_hash]
        
        # Compress image first
        print(f"[Nxdify] Compressing image (original: {original_size} bytes)...")
        compressed = self.compress_image_bytes_max(image_bytes, self.MAX_IMAGE_SIZE)
        compression_ratio = (1 - len(compressed) / original_size) * 100 if original_size > 0 else 0
        print(f"[Nxdify] Compressed to {len(compressed)} bytes ({compression_ratio:.1f}% reduction)")
        
        # Create temporary file
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
            tmp.write(compressed)
            tmp_path = tmp.name
        
        timeout_errors = []
        try:
            for attempt in range(max_attempts):
                try:
                    print(f"[Nxdify] Uploading image (attempt {attempt + 1}/{max_attempts})...")
                    attempt_start = time.time()
                    
                    # Upload to FAL (uses FAL_KEY environment variable set in process_async)
                    # upload_file expects a file path, not a BytesIO object
                    result = await asyncio.to_thread(self._upload_file_sync, tmp_path)
                    
                    attempt_elapsed = time.time() - attempt_start
                    print(f"[Nxdify] Upload completed in {attempt_elapsed:.2f} seconds")
                    
                    if isinstance(result, dict) and "url" in result:
                        url = result["url"]
                    elif isinstance(result, str):
                        url = result
                    else:
                        raise ValueError(f"Unexpected upload response: {result}")
                    
                    # Cache the URL if caching is enabled
                    if use_cache:
                        image_hash = self._compute_image_hash(image_bytes)
                        self._image_url_cache[image_hash] = url
                    
                    total_upload_time = time.time() - upload_start
                    print(f"[Nxdify] Image upload successful (total time: {total_upload_time:.2f} seconds)")
                    return url
                        
                except Exception as e:
                    # Check if this is a 408 Request Timeout error
                    error_str = str(e)
                    error_lower = error_str.lower()
                    is_408_timeout = (
                        "408" in error_str or
                        "request timeout" in error_lower or
                        "http/1.1 408" in error_lower or
                        "http 408" in error_lower
                    )
                    
                    # Check for other timeout errors
                    is_timeout = (
                        is_408_timeout or
                        "timeout" in error_lower or
                        isinstance(e, (TimeoutError, asyncio.TimeoutError)) or
                        (isinstance(e, aiohttp.ClientError) and "timeout" in error_lower)
                    )
                    
                    if is_408_timeout:
                        timeout_errors.append(f"Attempt {attempt + 1}: HTTP 408 Request Timeout")
                    
                    # If this is the last attempt and we had 408 timeouts, raise helpful exception
                    if attempt == max_attempts - 1:
                        if timeout_errors:
                            print(f"[Nxdify] Upload failed after {max_attempts} attempts")
                            raise RuntimeError(
                                f"Upload timed out after {max_attempts} attempts with HTTP 408 Request Timeout errors. "
                                f"The image may be too large. Please resize the image to a smaller resolution and try again. "
                                f"Errors: {'; '.join(timeout_errors)}"
                            )
                        print(f"[Nxdify] Upload failed on final attempt: {e}")
                        raise
                    
                    # If timeout error, retry with backoff
                    if is_timeout:
                        backoff = 2 + attempt * 3  # Exponential backoff: 2s, 5s, 8s
                        print(f"[Nxdify] Upload timeout error (attempt {attempt + 1}): {error_str[:100]}. Retrying in {backoff} seconds...")
                        await asyncio.sleep(backoff)
                        continue
                    
                    # Non-timeout error, fail immediately
                    print(f"[Nxdify] Upload failed with non-timeout error: {error_str[:100]}")
                    raise
        finally:
            # Clean up temp file
            try:
                os.unlink(tmp_path)
            except OSError:
                pass
    
    def _subscribe_sync(self, endpoint: str, arguments: dict):
        """Subscribe to FAL API job synchronously (handles submit + polling internally)."""
        print(f"[Nxdify] Submitting job to FAL API: {endpoint}")
        start_time = time.time()
        result = fal.subscribe(endpoint, arguments=arguments, with_logs=False)
        elapsed = time.time() - start_time
        print(f"[Nxdify] FAL API job completed in {elapsed:.2f} seconds")
        return result
    
    async def generate_image(
        self,
        face_url: str,
        body_url: str,
        breasts_url: str,
        dynamic_pose_url: str,
        prompt: str,
        quality: str
    ) -> Image.Image:
        """Generate image using FAL AI Seedream 4.5 API."""
        print(f"[Nxdify] Starting image generation with quality: {quality}")
        print(f"[Nxdify] Reference images: face={face_url[:50]}..., body={body_url[:50]}..., breasts={breasts_url[:50]}..., pose={dynamic_pose_url[:50]}...")
        
        image_urls = [face_url, body_url, breasts_url, dynamic_pose_url]
        
        arguments = {
            "prompt": prompt,
            "image_size": quality,
            "num_images": 1,
            "max_images": 1,
            "enable_safety_checker": False,
            "image_urls": image_urls
        }
        
        print(f"[Nxdify] Calling FAL API subscribe (this will poll internally)...")
        subscribe_start = time.time()
        
        # Use subscribe which handles submit + polling internally
        result = await asyncio.to_thread(
            self._subscribe_sync,
            "fal-ai/bytedance/seedream/v4.5/edit",
            arguments
        )
        
        subscribe_elapsed = time.time() - subscribe_start
        print(f"[Nxdify] Subscribe call returned after {subscribe_elapsed:.2f} seconds")
        
        if not result:
            raise ValueError("No result returned from FAL AI API")
        
        print(f"[Nxdify] Processing result (type: {type(result).__name__})...")
        
        # Extract images from result (handle different response structures)
        images = None
        if isinstance(result, dict):
            if "images" in result:
                images = result["images"]
                print(f"[Nxdify] Found {len(images)} image(s) in result['images']")
            elif "output" in result and isinstance(result["output"], dict):
                images = result["output"].get("images")
                print(f"[Nxdify] Found {len(images) if images else 0} image(s) in result['output']['images']")
        
        if not images or len(images) == 0:
            raise ValueError("No images returned from FAL AI API")
        
        # Handle both dict and string image URLs
        if isinstance(images[0], dict):
            image_url = images[0].get("url") or images[0].get("image_url")
        else:
            image_url = images[0]
        
        if not image_url:
            raise ValueError("No image URL in result")
        
        print(f"[Nxdify] Image URL extracted: {image_url[:80]}...")
        print(f"[Nxdify] Downloading generated image...")
        download_start = time.time()
        
        # Download image
        async with aiohttp.ClientSession() as session:
            async with session.get(image_url) as response:
                if response.status != 200:
                    raise ValueError(f"Failed to download image: HTTP {response.status}")
                image_bytes = await response.read()
        
        download_elapsed = time.time() - download_start
        print(f"[Nxdify] Image downloaded ({len(image_bytes)} bytes) in {download_elapsed:.2f} seconds")
        
        # Convert to PIL Image
        print(f"[Nxdify] Converting to PIL Image...")
        img = Image.open(io.BytesIO(image_bytes))
        final_img = img.convert("RGB")
        print(f"[Nxdify] Image generation complete. Final size: {final_img.size}")
        return final_img
    
    def pil_to_tensor(self, img: Image.Image) -> torch.Tensor:
        """Convert PIL Image to ComfyUI tensor format (BHWC)."""
        # Convert to RGB if needed
        if img.mode != "RGB":
            img = img.convert("RGB")
        
        # Convert to numpy array (HWC format: height, width, channels)
        img_array = np.array(img).astype(np.float32) / 255.0
        
        # Ensure shape is (height, width, channels)
        if len(img_array.shape) == 2:
            # Grayscale - add channel dimension
            img_array = np.expand_dims(img_array, axis=2)
            img_array = np.repeat(img_array, 3, axis=2)  # Convert to RGB
        
        # Add batch dimension to get BHWC format: (batch, height, width, channels)
        tensor = torch.from_numpy(img_array)[None,]
        return tensor
    
    async def process_async(
        self,
        face_image: torch.Tensor,
        body_image: torch.Tensor,
        breasts_image: torch.Tensor,
        dynamic_pose_image: torch.Tensor,
        prompt: str,
        fal_api_key: str,
        quality: str
    ) -> torch.Tensor:
        """Async processing function."""
        process_start = time.time()
        print(f"[Nxdify] ===== Starting Nxdify image generation process =====")
        
        if not fal_api_key:
            raise ValueError("FAL API key is required")
        
        if not prompt:
            raise ValueError("Prompt is required")
        
        print(f"[Nxdify] Converting input tensors to bytes...")
        # Convert tensors to bytes
        face_bytes = self.tensor_to_bytes(face_image)
        body_bytes = self.tensor_to_bytes(body_image)
        breasts_bytes = self.tensor_to_bytes(breasts_image)
        dynamic_pose_bytes = self.tensor_to_bytes(dynamic_pose_image)
        print(f"[Nxdify] Image sizes: face={len(face_bytes)} bytes, body={len(body_bytes)} bytes, breasts={len(breasts_bytes)} bytes, pose={len(dynamic_pose_bytes)} bytes")
        
        # Set FAL API key as environment variable (FAL SDK reads from env)
        os.environ["FAL_KEY"] = fal_api_key
        print(f"[Nxdify] FAL API key configured")
        
        print(f"[Nxdify] Uploading reference images...")
        upload_start = time.time()
        
        # Upload fixed references (cached - only upload if changed)
        print(f"[Nxdify] Uploading face image (cached)...")
        face_url = await self.upload_ref_with_retry(face_bytes, use_cache=True)
        print(f"[Nxdify] Uploading body image (cached)...")
        body_url = await self.upload_ref_with_retry(body_bytes, use_cache=True)
        print(f"[Nxdify] Uploading breasts image (cached)...")
        breasts_url = await self.upload_ref_with_retry(breasts_bytes, use_cache=True)
        
        # Upload dynamic pose image (not cached - always upload)
        print(f"[Nxdify] Uploading dynamic pose image (not cached)...")
        dynamic_pose_url = await self.upload_ref_with_retry(dynamic_pose_bytes, use_cache=False)
        
        upload_elapsed = time.time() - upload_start
        print(f"[Nxdify] All images uploaded in {upload_elapsed:.2f} seconds")
        
        # Generate image
        print(f"[Nxdify] Starting image generation...")
        generation_start = time.time()
        generated_img = await self.generate_image(
            face_url,
            body_url,
            breasts_url,
            dynamic_pose_url,
            prompt,
            quality
        )
        generation_elapsed = time.time() - generation_start
        print(f"[Nxdify] Image generation completed in {generation_elapsed:.2f} seconds")
        
        # Convert to tensor
        print(f"[Nxdify] Converting PIL image to tensor...")
        result = self.pil_to_tensor(generated_img)
        
        total_elapsed = time.time() - process_start
        print(f"[Nxdify] ===== Total process time: {total_elapsed:.2f} seconds =====")
        return result
    
    def execute(
        self,
        face_image: torch.Tensor,
        body_image: torch.Tensor,
        breasts_image: torch.Tensor,
        dynamic_pose_image: torch.Tensor,
        prompt: str,
        fal_api_key: str,
        quality: str
    ) -> Tuple[torch.Tensor]:
        """Execute the node (synchronous wrapper for async processing)."""
        # Handle event loop - check if one exists
        try:
            loop = asyncio.get_event_loop()
            if loop.is_running():
                # If loop is running, we need to use a different approach
                # Create a new event loop in a thread
                with concurrent.futures.ThreadPoolExecutor() as executor:
                    future = executor.submit(asyncio.run, self.process_async(
                        face_image,
                        body_image,
                        breasts_image,
                        dynamic_pose_image,
                        prompt,
                        fal_api_key,
                        quality
                    ))
                    result = future.result()
            else:
                # Loop exists but not running, use it
                result = loop.run_until_complete(self.process_async(
                    face_image,
                    body_image,
                    breasts_image,
                    dynamic_pose_image,
                    prompt,
                    fal_api_key,
                    quality
                ))
        except RuntimeError:
            # No event loop, create one
            result = asyncio.run(self.process_async(
                face_image,
                body_image,
                breasts_image,
                dynamic_pose_image,
                prompt,
                fal_api_key,
                quality
            ))
        
        return (result,)


# Node export
NODE_CLASS_MAPPINGS = {
    "NxdifyNode": NxdifyNode
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "NxdifyNode": "Nxdify Image Generation"
}