from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
import torch
import io
import base64
import tempfile
import os
from diffusers import AutoPipelineForInpainting, AutoencoderKL
from typing import Optional
import time

# Initialize FastAPI app
app = FastAPI(
    title="Virtual Try-On API",
    description="API for virtual clothing try-on using Stable Diffusion XL",
    version="1.0.0"
)

# Add CORS middleware for React Native
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Adjust for production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global variables for models
pipeline = None
segment_body = None

def load_models():
    """Load all required models"""
    global pipeline, segment_body
    
    # Determine device and dtype
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if torch.cuda.is_available() else torch.float32
    
    print(f"🔄 Using device: {device}, dtype: {dtype}")
    
    print("🔄 Loading VAE...")
    if torch.cuda.is_available():
        vae = AutoencoderKL.from_pretrained(
            "madebyollin/sdxl-vae-fp16-fix",
            torch_dtype=dtype
        )
    else:
        # Use standard VAE for CPU (fp32)
        vae = AutoencoderKL.from_pretrained(
            "stabilityai/sdxl-vae",
            torch_dtype=dtype
        )
    
    print("🔄 Loading inpainting pipeline...")
    if torch.cuda.is_available():
        pipeline = AutoPipelineForInpainting.from_pretrained(
            "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
            vae=vae,
            torch_dtype=dtype,
            variant="fp16",
            use_safetensors=True
        )
    else:
        # Load without fp16 variant for CPU
        pipeline = AutoPipelineForInpainting.from_pretrained(
            "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
            vae=vae,
            torch_dtype=dtype,
            use_safetensors=True
        )
    
    pipeline = pipeline.to(device)
    
    print("🔄 Loading IP-Adapter...")
    pipeline.load_ip_adapter(
        "h94/IP-Adapter",
        subfolder="sdxl_models",
        weight_name="ip-adapter_sdxl.bin",
        low_cpu_mem_usage=True
    )
    
    print("🔄 Loading body segmentation...")
    try:
        from SegBody import segment_body as seg_func
        segment_body = seg_func
        print("✅ Body segmentation loaded!")
    except ImportError:
        print("⚠️ SegBody module not found, segmentation will be disabled")
    
    print("✅ All models loaded successfully!")

@app.on_event("startup")
async def startup_event():
    """Load models on startup"""
    load_models()

@app.get("/")
async def root():
    """Health check endpoint"""
    return {
        "status": "running",
        "message": "Virtual Try-On API is running!",
        "cuda_available": torch.cuda.is_available(),
        "device": "cuda" if torch.cuda.is_available() else "cpu"
    }

@app.get("/health")
async def health():
    """Health check endpoint"""
    return {"status": "healthy", "models_loaded": pipeline is not None}

def image_to_base64(image: Image.Image) -> str:
    """Convert PIL Image to base64 string"""
    buffered = io.BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    return img_str

def base64_to_image(base64_str: str) -> Image.Image:
    """Convert base64 string to PIL Image"""
    img_data = base64.b64decode(base64_str)
    return Image.open(io.BytesIO(img_data)).convert('RGB')

@app.post("/tryon")
async def virtual_tryon(
    person_image: UploadFile = File(..., description="Image of the person"),
    clothing_image: UploadFile = File(..., description="Image of the clothing"),
    prompt: str = Form("photorealistic, perfect body, beautiful skin, realistic skin, natural skin"),
    negative_prompt: str = Form("ugly, bad quality, bad anatomy, deformed body, deformed hands, deformed feet, deformed face, deformed clothing, deformed skin, bad skin, leggings, tights, stockings"),
    ip_scale: float = Form(0.8),
    strength: float = Form(0.99),
    guidance_scale: float = Form(7.5),
    num_steps: int = Form(50),
    return_format: str = Form("base64", description="base64 or image")
):
    """
    Virtual Try-On endpoint
    
    Args:
        person_image: Image file of the person
        clothing_image: Image file of the clothing
        prompt: Generation prompt
        negative_prompt: Negative prompt
        ip_scale: IP-Adapter influence (0.0-1.0)
        strength: Inpainting strength (0.0-1.0)
        guidance_scale: CFG scale
        num_steps: Number of inference steps
        return_format: Response format (base64 or image)
    
    Returns:
        Generated image in specified format
    """
    try:
        if pipeline is None:
            raise HTTPException(status_code=503, detail="Models not loaded yet")
        
        start_time = time.time()
        
        # Load and resize images
        print("📥 Loading images...")
        person_img = Image.open(person_image.file).convert('RGB').resize((512, 512))
        clothing_img = Image.open(clothing_image.file).convert('RGB').resize((512, 512))
        
        # Generate body segmentation mask
        print("🎭 Generating segmentation mask...")
        if segment_body is None:
            # Create a simple fallback mask (full body) if segmentation not available
            mask_img = Image.new('L', (512, 512), 255)
        else:
            try:
                # Try calling segment_body - it might expect a file path or PIL Image
                try:
                    # First try with PIL Image directly
                    seg_image, mask_img = segment_body(person_img, face=False)
                except (TypeError, AttributeError):
                    # If that fails, try with file path
                    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
                        temp_path = tmp_file.name
                        person_img.save(temp_path)
                    
                    seg_image, mask_img = segment_body(temp_path, face=False)
                    os.unlink(temp_path)
                
                # Ensure mask is PIL Image and resize
                if isinstance(mask_img, str):
                    mask_img = Image.open(mask_img).convert('L')
                mask_img = mask_img.resize((512, 512))
                
            except Exception as e:
                print(f"⚠️ Segmentation failed: {e}, using full mask")
                mask_img = Image.new('L', (512, 512), 255)
        
        # Set IP-Adapter scale
        pipeline.set_ip_adapter_scale(ip_scale)
        
        # Generate virtual try-on
        print("🎨 Generating virtual try-on...")
        result = pipeline(
            prompt=prompt,
            negative_prompt=negative_prompt,
            image=person_img,
            mask_image=mask_img,
            ip_adapter_image=clothing_img,
            strength=strength,
            guidance_scale=guidance_scale,
            num_inference_steps=num_steps,
        )
        
        generated_image = result.images[0]
        
        processing_time = time.time() - start_time
        print(f"✅ Generation completed in {processing_time:.2f}s")
        
        # Return based on format
        if return_format == "image":
            # Return as image file
            img_byte_arr = io.BytesIO()
            generated_image.save(img_byte_arr, format='PNG')
            img_byte_arr.seek(0)
            return StreamingResponse(img_byte_arr, media_type="image/png")
        else:
            # Return as base64 JSON
            img_base64 = image_to_base64(generated_image)
            return JSONResponse({
                "success": True,
                "image": img_base64,
                "processing_time": processing_time,
                "parameters": {
                    "prompt": prompt,
                    "ip_scale": ip_scale,
                    "strength": strength,
                    "guidance_scale": guidance_scale,
                    "num_steps": num_steps
                }
            })
    
    except Exception as e:
        print(f"❌ Error: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/tryon-base64")
async def virtual_tryon_base64(
    person_image_base64: str = Form(..., description="Base64 encoded person image"),
    clothing_image_base64: str = Form(..., description="Base64 encoded clothing image"),
    prompt: str = Form("photorealistic, perfect body, beautiful skin, realistic skin, natural skin"),
    negative_prompt: str = Form("ugly, bad quality, bad anatomy, deformed body, deformed hands, deformed feet, deformed face, deformed clothing, deformed skin, bad skin, leggings, tights, stockings"),
    ip_scale: float = Form(0.8),
    strength: float = Form(0.99),
    guidance_scale: float = Form(7.5),
    num_steps: int = Form(50)
):
    """
    Virtual Try-On endpoint accepting base64 encoded images
    (Alternative endpoint for easier React Native integration)
    """
    try:
        if pipeline is None:
            raise HTTPException(status_code=503, detail="Models not loaded yet")
        
        start_time = time.time()
        
        # Decode base64 images
        print("📥 Decoding base64 images...")
        person_img = base64_to_image(person_image_base64).resize((512, 512))
        clothing_img = base64_to_image(clothing_image_base64).resize((512, 512))
        
        # Generate body segmentation mask
        print("🎭 Generating segmentation mask...")
        if segment_body is None:
            mask_img = Image.new('L', (512, 512), 255)
        else:
            try:
                # Try calling segment_body - it might expect a file path or PIL Image
                try:
                    # First try with PIL Image directly
                    seg_image, mask_img = segment_body(person_img, face=False)
                except (TypeError, AttributeError):
                    # If that fails, try with file path
                    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
                        temp_path = tmp_file.name
                        person_img.save(temp_path)
                    
                    seg_image, mask_img = segment_body(temp_path, face=False)
                    os.unlink(temp_path)
                
                # Ensure mask is PIL Image and resize
                if isinstance(mask_img, str):
                    mask_img = Image.open(mask_img).convert('L')
                mask_img = mask_img.resize((512, 512))
                
            except Exception as e:
                print(f"⚠️ Segmentation failed: {e}, using full mask")
                mask_img = Image.new('L', (512, 512), 255)
        
        # Set IP-Adapter scale
        pipeline.set_ip_adapter_scale(ip_scale)
        
        # Generate virtual try-on
        print("🎨 Generating virtual try-on...")
        result = pipeline(
            prompt=prompt,
            negative_prompt=negative_prompt,
            image=person_img,
            mask_image=mask_img,
            ip_adapter_image=clothing_img,
            strength=strength,
            guidance_scale=guidance_scale,
            num_inference_steps=num_steps,
        )
        
        generated_image = result.images[0]
        processing_time = time.time() - start_time
        print(f"✅ Generation completed in {processing_time:.2f}s")
        
        # Return as base64
        img_base64 = image_to_base64(generated_image)
        return JSONResponse({
            "success": True,
            "image": img_base64,
            "processing_time": processing_time,
            "parameters": {
                "prompt": prompt,
                "ip_scale": ip_scale,
                "strength": strength,
                "guidance_scale": guidance_scale,
                "num_steps": num_steps
            }
        })
    
    except Exception as e:
        print(f"❌ Error: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)