import os
import math
import gradio as gr
import numpy as np
import torch
import safetensors.torch as sf
from datetime import datetime

# Import spaces for GPU decorator
try:
    import spaces
    HF_SPACES_GPU = True
except ImportError:
    HF_SPACES_GPU = False
    # Create a dummy decorator if spaces is not available
    class spaces:
        @staticmethod
        def GPU(func):
            return func

from PIL import Image
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
from diffusers import AutoencoderKL, UNet2DConditionModel, DDIMScheduler, EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler
from transformers import CLIPTextModel, CLIPTokenizer
from enum import Enum
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import PyTorchModelHubMixin

# Try to import RMBG, fallback to local implementation
try:
    from transformers import pipeline
    rmbg_pipeline = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
    USE_RMBG_PIPELINE = True
except Exception as e:
    print(f"Failed to load RMBG pipeline: {e}")
    USE_RMBG_PIPELINE = False
    try:
        from briarmbg import BriaRMBG, simple_background_removal
    except:
        # Inline simple background removal
        def simple_background_removal(image):
            if isinstance(image, np.ndarray):
                img = image
            else:
                img = np.array(image)
            
            # Simple fallback - return full mask
            gray = np.mean(img, axis=2)
            mask = np.ones_like(gray)
            return mask

# Model setup
sd15_name = 'stablediffusionapi/realistic-vision-v51'

# Better CUDA detection and debugging
print("===== Application Startup at", datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "=====")
print()
print("=== GPU Detection Debug ===")
print(f"PyTorch version: {torch.__version__}")
print(f"Hugging Face Spaces GPU support: {HF_SPACES_GPU}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"CUDA version: {torch.version.cuda}")
    print(f"GPU count: {torch.cuda.device_count()}")
    print(f"Current GPU: {torch.cuda.current_device()}")
    print(f"GPU name: {torch.cuda.get_device_name()}")
    print("✅ GPU detected and available!")
else:
    print("❌ CUDA not available - checking reasons...")
    try:
        import subprocess
        result = subprocess.run(['nvidia-smi'], capture_output=True, text=True)
        if result.returncode == 0:
            print("nvidia-smi works, GPU hardware detected")
            print("Issue might be with PyTorch CUDA installation")
        else:
            print("nvidia-smi failed, no GPU hardware detected")
    except:
        print("nvidia-smi command not found")
    
    if HF_SPACES_GPU:
        print("🔄 Running on Hugging Face Spaces with @spaces.GPU decorator")
        print("   GPU will be allocated when GPU-decorated functions are called")
    else:
        print()
        print("🚨 WARNING: This application requires GPU to run properly!")
        print("📋 To fix this issue:")
        print("   1. Go to your Space settings: https://huggingface.co/spaces/GreenGoat/IClight-demo/settings")
        print("   2. In the Hardware section, select 'GPU basic' or higher")
        print("   3. Make sure your Hugging Face account is verified")
        print("   4. Check if you have available GPU quota")
    print()

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Selected device: {device}")
print("=== End GPU Debug ===")

print(f"Using device: {device}")
print("Loading models...")

# Initialize models
tokenizer = CLIPTokenizer.from_pretrained(sd15_name, subfolder="tokenizer")
text_encoder = CLIPTextModel.from_pretrained(sd15_name, subfolder="text_encoder")
vae = AutoencoderKL.from_pretrained(sd15_name, subfolder="vae")
unet = UNet2DConditionModel.from_pretrained(sd15_name, subfolder="unet")

# Modify UNet for IC-Light
with torch.no_grad():
    new_conv_in = torch.nn.Conv2d(12, unet.conv_in.out_channels, unet.conv_in.kernel_size, unet.conv_in.stride, unet.conv_in.padding)
    new_conv_in.weight.zero_()
    new_conv_in.weight[:, :4, :, :].copy_(unet.conv_in.weight)
    new_conv_in.bias = unet.conv_in.bias
    unet.conv_in = new_conv_in

unet_original_forward = unet.forward

def hooked_unet_forward(sample, timestep, encoder_hidden_states, **kwargs):
    c_concat = kwargs['cross_attention_kwargs']['concat_conds'].to(sample)
    c_concat = torch.cat([c_concat] * (sample.shape[0] // c_concat.shape[0]), dim=0)
    new_sample = torch.cat([sample, c_concat], dim=1)
    kwargs['cross_attention_kwargs'] = {}
    return unet_original_forward(new_sample, timestep, encoder_hidden_states, **kwargs)

unet.forward = hooked_unet_forward

# Load IC-Light weights
model_path = './iclight_sd15_fbc.safetensors'
if not os.path.exists(model_path):
    print("Downloading IC-Light model...")
    try:
        from huggingface_hub import hf_hub_download
        model_path = hf_hub_download(
            repo_id="lllyasviel/ic-light",
            filename="iclight_sd15_fbc.safetensors"
        )
    except Exception as e:
        print(f"Failed to download with hf_hub_download: {e}")
        # Fallback to torch.hub
        from torch.hub import download_url_to_file
        download_url_to_file(url='https://huggingface.co/lllyasviel/ic-light/resolve/main/iclight_sd15_fbc.safetensors', dst=model_path)

sd_offset = sf.load_file(model_path)
sd_origin = unet.state_dict()
sd_merged = {k: sd_origin[k] + sd_offset[k] for k in sd_origin.keys()}
unet.load_state_dict(sd_merged, strict=True)
del sd_offset, sd_origin, sd_merged

# Move models to device
text_encoder = text_encoder.to(device=device, dtype=torch.float16)
vae = vae.to(device=device, dtype=torch.bfloat16)
unet = unet.to(device=device, dtype=torch.float16)

# Scheduler
scheduler = DPMSolverMultistepScheduler(
    num_train_timesteps=1000,
    beta_start=0.00085,
    beta_end=0.012,
    algorithm_type="sde-dpmsolver++",
    use_karras_sigmas=True,
    steps_offset=1
)

# Pipelines
t2i_pipe = StableDiffusionPipeline(
    vae=vae,
    text_encoder=text_encoder,
    tokenizer=tokenizer,
    unet=unet,
    scheduler=scheduler,
    safety_checker=None,
    requires_safety_checker=False,
    feature_extractor=None
)

i2i_pipe = StableDiffusionImg2ImgPipeline(
    vae=vae,
    text_encoder=text_encoder,
    tokenizer=tokenizer,
    unet=unet,
    scheduler=scheduler,
    safety_checker=None,
    requires_safety_checker=False,
    feature_extractor=None
)

print("Models loaded successfully!")

@torch.inference_mode()
def encode_prompt_inner(txt: str):
    max_length = tokenizer.model_max_length
    chunk_length = tokenizer.model_max_length - 2
    id_start = tokenizer.bos_token_id
    id_end = tokenizer.eos_token_id
    id_pad = id_end

    def pad(x, p, i):
        return x[:i] if len(x) >= i else x + [p] * (i - len(x))

    tokens = tokenizer(txt, truncation=False, add_special_tokens=False)["input_ids"]
    chunks = [[id_start] + tokens[i: i + chunk_length] + [id_end] for i in range(0, len(tokens), chunk_length)]
    chunks = [pad(ck, id_pad, max_length) for ck in chunks]

    token_ids = torch.tensor(chunks).to(device=device, dtype=torch.int64)
    conds = text_encoder(token_ids).last_hidden_state

    return conds

@torch.inference_mode()
def encode_prompt_pair(positive_prompt, negative_prompt):
    c = encode_prompt_inner(positive_prompt)
    uc = encode_prompt_inner(negative_prompt)

    c_len = float(len(c))
    uc_len = float(len(uc))
    max_count = max(c_len, uc_len)
    c_repeat = int(math.ceil(max_count / c_len))
    uc_repeat = int(math.ceil(max_count / uc_len))
    max_chunk = max(len(c), len(uc))

    c = torch.cat([c] * c_repeat, dim=0)[:max_chunk]
    uc = torch.cat([uc] * uc_repeat, dim=0)[:max_chunk]

    c = torch.cat([p[None, ...] for p in c], dim=1)
    uc = torch.cat([p[None, ...] for p in uc], dim=1)

    return c, uc

@torch.inference_mode()
def pytorch2numpy(imgs, quant=True):
    results = []
    for x in imgs:
        y = x.movedim(0, -1)
        if quant:
            y = y * 127.5 + 127.5
            y = y.detach().float().cpu().numpy().clip(0, 255).astype(np.uint8)
        else:
            y = y * 0.5 + 0.5
            y = y.detach().float().cpu().numpy().clip(0, 1).astype(np.float32)
        results.append(y)
    return results

@torch.inference_mode()
def numpy2pytorch(imgs):
    h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.0 - 1.0
    h = h.movedim(-1, 1)
    return h

def resize_and_center_crop(image, target_width, target_height):
    pil_image = Image.fromarray(image)
    original_width, original_height = pil_image.size
    scale_factor = max(target_width / original_width, target_height / original_height)
    new_width = int(original_width * scale_factor)
    new_height = int(original_height * scale_factor)
    pil_image = pil_image.resize((new_width, new_height), Image.LANCZOS)
    left = (new_width - target_width) / 2
    top = (new_height - target_height) / 2
    right = (new_width + target_width) / 2
    bottom = (new_height + target_height) / 2
    pil_image = pil_image.crop((left, top, right, bottom))
    return np.array(pil_image)

def resize_without_crop(image, target_width, target_height):
    pil_image = Image.fromarray(image)
    pil_image = pil_image.resize((target_width, target_height), Image.LANCZOS)
    return np.array(pil_image)

@spaces.GPU
@torch.inference_mode()
def run_rmbg(img, sigma=0.0):
    # Simplified background removal
    if USE_RMBG_PIPELINE:
        # Using transformers pipeline
        try:
            result = rmbg_pipeline(Image.fromarray(img))
            mask = np.array(result['mask'])
            if len(mask.shape) == 3:
                mask = mask[:, :, 0]
            mask = mask.astype(np.float32) / 255.0
        except Exception as e:
            print(f"RMBG pipeline failed: {e}, using fallback")
            mask = simple_background_removal(img)
    else:
        # Using simple background removal
        mask = simple_background_removal(img)
    
    # Apply sigma smoothing
    if sigma > 0:
        try:
            from scipy import ndimage
            mask = ndimage.gaussian_filter(mask, sigma=sigma)
        except ImportError:
            # Fallback if scipy is not available
            pass
    
    # Create RGBA output
    result = np.dstack((img, (mask * 255).astype(np.uint8)))
    return img, mask

class BGSource(Enum):
    UPLOAD = "Use Background Image"
    UPLOAD_FLIP = "Use Flipped Background Image"
    LEFT = "Left Light"
    RIGHT = "Right Light"
    TOP = "Top Light"
    BOTTOM = "Bottom Light"
    GREY = "Ambient"

@spaces.GPU
@torch.inference_mode()
def process(input_fg, input_bg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, bg_source):
    bg_source = BGSource(bg_source)

    if bg_source == BGSource.UPLOAD:
        pass
    elif bg_source == BGSource.UPLOAD_FLIP:
        input_bg = np.fliplr(input_bg)
    elif bg_source == BGSource.GREY:
        input_bg = np.zeros(shape=(image_height, image_width, 3), dtype=np.uint8) + 64
    elif bg_source == BGSource.LEFT:
        gradient = np.linspace(224, 32, image_width)
        image = np.tile(gradient, (image_height, 1))
        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
    elif bg_source == BGSource.RIGHT:
        gradient = np.linspace(32, 224, image_width)
        image = np.tile(gradient, (image_height, 1))
        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
    elif bg_source == BGSource.TOP:
        gradient = np.linspace(224, 32, image_height)[:, None]
        image = np.tile(gradient, (1, image_width))
        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
    elif bg_source == BGSource.BOTTOM:
        gradient = np.linspace(32, 224, image_height)[:, None]
        image = np.tile(gradient, (1, image_width))
        input_bg = np.stack((image,) * 3, axis=-1).astype(np.uint8)
    else:
        raise ValueError('Wrong background source!')

    rng = torch.Generator(device=device).manual_seed(seed)

    fg = resize_and_center_crop(input_fg, image_width, image_height)
    bg = resize_and_center_crop(input_bg, image_width, image_height)
    concat_conds = numpy2pytorch([fg, bg]).to(device=vae.device, dtype=vae.dtype)
    concat_conds = vae.encode(concat_conds).latent_dist.mode() * vae.config.scaling_factor
    concat_conds = torch.cat([c[None, ...] for c in concat_conds], dim=1)

    conds, unconds = encode_prompt_pair(positive_prompt=prompt + ', ' + a_prompt, negative_prompt=n_prompt)

    latents = t2i_pipe(
        prompt_embeds=conds,
        negative_prompt_embeds=unconds,
        width=image_width,
        height=image_height,
        num_inference_steps=steps,
        num_images_per_prompt=num_samples,
        generator=rng,
        output_type='latent',
        guidance_scale=cfg,
        cross_attention_kwargs={'concat_conds': concat_conds},
    ).images.to(vae.dtype) / vae.config.scaling_factor

    pixels = vae.decode(latents).sample
    pixels = pytorch2numpy(pixels)  # Use default quant=True for first pass
    
    # Always perform highres processing like the original code
    pixels = [resize_without_crop(
        image=p,
        target_width=int(round(image_width * highres_scale / 64.0) * 64),
        target_height=int(round(image_height * highres_scale / 64.0) * 64))
    for p in pixels]

    pixels = numpy2pytorch(pixels).to(device=vae.device, dtype=vae.dtype)
    latents = vae.encode(pixels).latent_dist.mode() * vae.config.scaling_factor
    latents = latents.to(device=unet.device, dtype=unet.dtype)

    image_height, image_width = latents.shape[2] * 8, latents.shape[3] * 8
    fg = resize_and_center_crop(input_fg, image_width, image_height)
    bg = resize_and_center_crop(input_bg, image_width, image_height)
    concat_conds = numpy2pytorch([fg, bg]).to(device=vae.device, dtype=vae.dtype)
    concat_conds = vae.encode(concat_conds).latent_dist.mode() * vae.config.scaling_factor
    concat_conds = torch.cat([c[None, ...] for c in concat_conds], dim=1)

    latents = i2i_pipe(
        image=latents,
        strength=highres_denoise,
        prompt_embeds=conds,
        negative_prompt_embeds=unconds,
        width=image_width,
        height=image_height,
        num_inference_steps=int(round(steps / highres_denoise)),
        num_images_per_prompt=num_samples,
        generator=rng,
        output_type='latent',
        guidance_scale=cfg,
        cross_attention_kwargs={'concat_conds': concat_conds},
    ).images.to(vae.dtype) / vae.config.scaling_factor

    pixels = vae.decode(latents).sample
    pixels = pytorch2numpy(pixels, quant=False)  # Return 0-1 range floats for final result

    return pixels, [fg, bg]

@spaces.GPU
@torch.inference_mode()
def process_relight(input_fg, input_bg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, bg_source):
    try:
        # Input validation
        if input_fg is None:
            error_msg = "❌ Please upload a foreground image"
            print(error_msg)
            raise gr.Error(error_msg)
        if input_bg is None and bg_source == "Use Background Image":
            error_msg = "❌ Please upload a background image or choose a lighting direction"
            print(error_msg)
            raise gr.Error(error_msg)
        
        # Handle empty prompt - provide default when using background image
        if not prompt.strip():
            if bg_source == "Use Background Image" or bg_source == "Use Flipped Background Image":
                # When using background image as light source, use a generic default prompt
                prompt = "best quality, detailed"
                print(f"Using default prompt for background lighting: {prompt}")
            else:
                error_msg = "❌ Please enter a prompt"
                print(error_msg)
                raise gr.Error(error_msg)
        
        print(f"Processing with device: {device}")
        print(f"Input shapes - FG: {input_fg.shape}, BG: {input_bg.shape if input_bg is not None else 'None'}")
        
        # Optimize for Hugging Face free GPU (limited memory)
        if device.type == 'cuda':
            # Limit image size for free GPU tier
            max_size = 768  # Increased for GPU but still conservative
            if image_width > max_size or image_height > max_size:
                scale = min(max_size / image_width, max_size / image_height)
                image_width = int(image_width * scale // 64) * 64  # Keep multiple of 64
                image_height = int(image_height * scale // 64) * 64
                print(f"Reduced image size for GPU memory: {image_width}x{image_height}")
            
            # Disable highres for free tier to save memory
            if highres_scale > 1.0:
                highres_scale = 1.0
                print("Disabled highres scaling to save GPU memory")
                
        elif device.type == 'cpu':
            # Limit image size for CPU processing
            max_size = 512
            if image_width > max_size or image_height > max_size:
                image_width = min(image_width, max_size)
                image_height = min(image_height, max_size)
                print(f"Reduced image size for CPU: {image_width}x{image_height}")
            
            # Limit number of samples for CPU
            if num_samples > 1:
                num_samples = 1
                print("Reduced num_samples to 1 for CPU processing")
        
        print("Running background removal...")
        try:
            input_fg, matting = run_rmbg(input_fg)
            print("Background removal completed successfully")
        except Exception as e:
            print(f"Background removal failed: {e}")
            # Continue without background removal
            matting = np.ones((input_fg.shape[0], input_fg.shape[1]), dtype=np.float32)
        
        print("Starting main processing...")
        try:
            results, extra_images = process(input_fg, input_bg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, bg_source)
            print("Main processing completed successfully")
        except Exception as e:
            error_msg = f"❌ Processing failed: {str(e)}"
            print(error_msg)
            import traceback
            traceback.print_exc()
            raise gr.Error(error_msg)
        
        print("Converting results...")
        try:
            results = [(x * 255.0).clip(0, 255).astype(np.uint8) for x in results]
            print("Results converted successfully")
        except Exception as e:
            error_msg = f"❌ Result conversion failed: {str(e)}"
            print(error_msg)
            raise gr.Error(error_msg)
        
        print("Processing completed successfully!")
        return results + extra_images
        
    except gr.Error:
        # Re-raise Gradio errors to show them in the UI
        raise
    except Exception as e:
        error_msg = f"❌ Unexpected error: {str(e)}"
        print(error_msg)
        import traceback
        traceback.print_exc()
        raise gr.Error(error_msg)

# Quick prompts for easy testing
quick_prompts = [
    'beautiful woman, cinematic lighting',
    'handsome man, cinematic lighting',
    'beautiful woman, natural lighting',
    'handsome man, natural lighting',
    'beautiful woman, neo punk lighting, cyberpunk',
    'handsome man, neo punk lighting, cyberpunk',
]
quick_prompts = [[x] for x in quick_prompts]

# Gradio Interface
def create_demo():
    with gr.Blocks(title="IC-Light Background Conditional Relighting") as demo:
        gr.Markdown("## IC-Light: Relighting with Foreground and Background Condition")
        gr.Markdown("Upload a foreground image and background image (or choose lighting direction) to perform relighting.")
        
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    input_fg = gr.Image(label="Foreground Image", height=400, type="numpy")
                    input_bg = gr.Image(label="Background Image", height=400, type="numpy")
                
                prompt = gr.Textbox(label="Prompt", value="beautiful woman, cinematic lighting")
                bg_source = gr.Radio(
                    choices=[e.value for e in BGSource],
                    value=BGSource.UPLOAD.value,
                    label="Background Source"
                )
                
                example_prompts = gr.Dataset(
                    samples=quick_prompts, 
                    label='Quick Prompts', 
                    components=[prompt]
                )
                
                relight_button = gr.Button(value="✨ Relight Image", variant="primary")
                
                with gr.Accordion("Advanced Settings", open=False):
                    with gr.Row():
                        num_samples = gr.Slider(label="Number of Images", minimum=1, maximum=4, value=1, step=1)
                        seed = gr.Number(label="Seed", value=12345, precision=0)
                    with gr.Row():
                        image_width = gr.Slider(label="Width", minimum=256, maximum=1024, value=512, step=64)
                        image_height = gr.Slider(label="Height", minimum=256, maximum=1024, value=640, step=64)
                    with gr.Row():
                        steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=20, step=1)
                        cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=20.0, value=7.0, step=0.1)
                    with gr.Row():
                        highres_scale = gr.Slider(label="Highres Scale", minimum=1.0, maximum=2.0, value=1.5, step=0.1)
                        highres_denoise = gr.Slider(label="Highres Denoise", minimum=0.1, maximum=0.9, value=0.5, step=0.1)
                    
                    a_prompt = gr.Textbox(label="Additional Prompt", value='best quality')
                    n_prompt = gr.Textbox(label="Negative Prompt", value='lowres, bad anatomy, bad hands, cropped, worst quality')
            
            with gr.Column():
                result_gallery = gr.Gallery(label='Results', height=600, columns=2, rows=2)
        
        # Event handlers
        inputs = [input_fg, input_bg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, bg_source]
        relight_button.click(
            fn=process_relight, 
            inputs=inputs, 
            outputs=[result_gallery],
            show_progress=True
        )
        example_prompts.click(lambda x: x[0], inputs=example_prompts, outputs=prompt, show_progress=False)
        
        # Examples - temporarily disabled due to missing image files
        # gr.Examples(
        #     examples=[
        #         ["examples/person1.jpg", "examples/bg1.jpg", "beautiful woman, cinematic lighting", "Use Background Image"],
        #         ["examples/person2.jpg", None, "handsome man, dramatic lighting", "Left Light"],
        #     ],
        #     inputs=[input_fg, input_bg, prompt, bg_source],
        #     outputs=[result_gallery],
        #     fn=process_relight,
        #     cache_examples=False,
        # )
    
    return demo

if __name__ == "__main__":
    demo = create_demo()
    demo.queue(max_size=20)
    demo.launch(
        server_name='0.0.0.0', 
        server_port=7860,
        show_error=True,
        share=False
    )