"""
Super Resolution Anime Diffusion - Enhanced WebUI

This is an enhanced version of the original Super Resolution Anime Diffusion project by yangheng95.
The WebUI has been improved with modern Gradio API implementation, better user experience,
and comprehensive documentation.

Key Contributions:
- Updated to use modern Gradio Blocks API for better interface organization
- Added tabbed interface for Text-to-Image, Image-to-Image, and Gallery views
- Improved error handling and user feedback with progress indicators
- Enhanced UI styling with custom CSS and responsive design
- Better parameter organization with collapsible accordions
- Real-time system information display

Instructions:
1. Choose between Text-to-Image or Image-to-Image tabs
2. Select a model from the dropdown (or provide custom model path)
3. Enter your prompt and adjust parameters as needed
4. For Image-to-Image: upload a base image to transform
5. Configure super-resolution settings (method and scale factor)
6. Click Generate to create high-quality anime images with automatic upscaling

Original Author: yangheng95
Original Repository: https://github.com/yangheng95/SuperResolutionAnimeDiffusion
License: Creative ML Open RAIL-M
Enhanced WebUI by AI Assistant
"""

import os
import sys
import zipfile
from typing import Optional, List, Tuple
from datetime import datetime
import time
import psutil

import PIL.Image
import autocuda
import findfile

from diffusers import (
    AutoencoderKL,
    UNet2DConditionModel,
    StableDiffusionPipeline,
    StableDiffusionImg2ImgPipeline,
    DPMSolverMultistepScheduler,
)
import gradio as gr
import torch
from PIL import Image
import utils
from Waifu2x.magnify import ImageMagnifier
from RealESRGANv030.interface import realEsrgan

sys.path.append(os.path.dirname(__file__))  # Ensure current directory is in path
os.environ["PYTHONPATH"] = os.path.dirname(__file__)

# Application Configuration
APP_TITLE = "🎨 Super Resolution Anime Diffusion"
APP_DESCRIPTION = """
Generate high-quality anime images with automatic super resolution enhancement.
Combines Stable Diffusion models with advanced upscaling techniques (RealESRGAN & Waifu2x).
"""

CONTRIBUTION_INFO = """
### 🤝 Enhanced Features
This interface improves upon the original work with:
- **Modern UI**: Clean tabbed interface with Gradio Blocks
- **Better UX**: Progress tracking and real-time feedback
- **Enhanced Parameters**: Organized controls with descriptions
- **Gallery View**: Browse and manage generated images
- **Error Handling**: Comprehensive error reporting and recovery
"""

INSTRUCTIONS = """
### 🚀 How to Use
1. **Select Mode**: Choose Text-to-Image or Image-to-Image tab
2. **Pick Model**: Select from available models or use custom path
3. **Create Prompt**: Describe your desired image (use negative prompt to avoid elements)
4. **Upload Image**: For img2img mode, provide base image
5. **Adjust Settings**: Fine-tune resolution, steps, and guidance
6. **Set Upscaling**: Choose super-resolution method and scale
7. **Generate**: Click the generate button and wait for results!
"""

COPYRIGHT_INFO = """
**Original Author**: [yangheng95](https://github.com/yangheng95) | 
**Repository**: [SuperResolutionAnimeDiffusion](https://github.com/yangheng95/SuperResolutionAnimeDiffusion) | 
**License**: Creative ML Open RAIL-M | **Enhanced by**: AI Assistant
"""

DEFAULT_NEGATIVE_PROMPT = "bad result, worst, random, invalid, inaccurate, imperfect, blurry, deformed, disfigured, mutation, mutated, ugly, out of focus, bad anatomy, text, error, extra digit, fewer digits, worst quality, low quality, normal quality, noise, jpeg artifact, compression artifact, signature, watermark, username, logo, low resolution, worst resolution, bad resolution, normal resolution, bad detail, bad details, bad lighting, bad shadow, bad shading, bad background, worst background"

# Initialization
magnifier = ImageMagnifier()
start_time = time.time()
is_colab = utils.is_google_colab()
device = autocuda.auto_cuda()
dtype = torch.float16 if device != "cpu" else torch.float32

# Extract zip files if needed
for z_file in findfile.find_cwd_files(and_key=['.zip'], exclude_key=['.ignore'], recursive=1):
    try:
        with zipfile.ZipFile(z_file, 'r') as zip_ref:
            zip_ref.extractall()
    except Exception as e:
        print(f"Warning: Could not extract {z_file}: {e}")

class Model:
    """Model configuration class"""
    def __init__(self, name: str, path: str = "", prefix: str = ""):
        self.name = name
        self.path = path
        self.prefix = prefix
        self.pipe_t2i = None
        self.pipe_i2i = None

# Model configurations
models = [
    Model("Anything v4.5", "xyn-ai/anything-v4.0", "anything v4.5 style"),
]

# Scheduler configuration
scheduler = DPMSolverMultistepScheduler.from_config({
    "beta_start": 0.00085,
    "beta_end": 0.012,
    "beta_schedule": "scaled_linear",
    "num_train_timesteps": 1000,
    "trained_betas": None,
    "prediction_type": "epsilon",
    "thresholding": False,
    "algorithm_type": "dpmsolver++",
    "solver_type": "midpoint",
    "solver_order": 2,
    "use_karras_sigmas": False,
    "timestep_spacing": "leading",
    "steps_offset": 1
})

# Global state
custom_model = None
if is_colab:
    models.insert(0, Model("Custom model"))
    custom_model = models[0]

last_mode = "txt2img"
current_model = models[1] if is_colab else models[0]
current_model_path = current_model.path
pipe = None

def initialize_models():
    """Initialize diffusion models with error handling"""
    global pipe

    if is_colab:
        try:
            pipe = StableDiffusionPipeline.from_pretrained(
                current_model.path,
                torch_dtype=dtype,
                scheduler=scheduler,
                safety_checker=None,
            )
        except Exception as e:
            print(f"Failed to initialize model: {e}")
            return
    else:
        print(f"{datetime.now()} Loading models...")
        try:
            vae = AutoencoderKL.from_pretrained(
                current_model.path, subfolder="vae", torch_dtype=dtype
            )

            for model in models[:]:
                try:
                    print(f"Loading {model.name}...")
                    unet = UNet2DConditionModel.from_pretrained(
                        model.path, subfolder="unet", torch_dtype=dtype
                    )
                    model.pipe_t2i = StableDiffusionPipeline.from_pretrained(
                        model.path,
                        unet=unet,
                        vae=vae,
                        torch_dtype=dtype,
                        scheduler=scheduler,
                        safety_checker=None,
                    )
                    model.pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(
                        model.path,
                        unet=unet,
                        vae=vae,
                        torch_dtype=dtype,
                        scheduler=scheduler,
                        safety_checker=None,
                    )
                    print(f"✅ {model.name} loaded successfully")
                except Exception as e:
                    print(f"❌ Failed to load {model.name}: {e}")
                    models.remove(model)

            if models:
                pipe = models[0].pipe_t2i
        except Exception as e:
            print(f"Failed to initialize models: {e}")
            return

    if torch.cuda.is_available() and pipe:
        pipe = pipe.to(device)

def get_system_info() -> str:
    """Get system information"""
    gpu_name = "CPU"
    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name()

    memory = psutil.virtual_memory()
    return f"🖥️ Device: {gpu_name} | 💾 RAM: {memory.available // (1024**3):.1f}GB"

def error_str(error: Exception, title: str = "Error") -> str:
    """Format error messages"""
    return f"### ❌ {title}\n```\n{str(error)}\n```"

def custom_model_changed(path: str) -> str:
    """Handle custom model path changes"""
    if custom_model and path.strip():
        models[0].path = path.strip()
        global current_model
        current_model = models[0]
        return "✅ Custom model path updated"
    return "❌ Please enter a valid model path"

def on_model_change(model_name: str) -> Tuple[gr.update, gr.update]:
    """Handle model selection changes"""
    selected_model = next((m for m in models if m.name == model_name), None)

    if selected_model and selected_model != models[0] if custom_model else True:
        prefix_text = f'Prompt (automatically prefixed with "{selected_model.prefix}")'
        is_custom = False
    else:
        prefix_text = "Enter prompt (remember to include model-specific prefix)"
        is_custom = True

    return (
        gr.update(visible=is_custom),
        gr.update(placeholder=prefix_text),
    )

def generate_image(
    mode: str,
    model_name: str,
    prompt: str,
    negative_prompt: str,
    width: int,
    height: int,
    guidance_scale: float,
    num_steps: int,
    seed: int,
    image: Optional[PIL.Image.Image],
    strength: float,
    scale_method: str,
    scale_factor: int,
    progress=gr.Progress()
) -> Tuple[Optional[PIL.Image.Image], str]:
    """Main image generation function"""

    if progress:
        progress(0, desc="Starting generation...")

    # Validation
    if not prompt.strip():
        return None, "❌ Please enter a prompt"

    if mode == "img2img" and image is None:
        return None, "❌ Please upload an image for Image-to-Image mode"

    # Find model
    global current_model
    selected_model = next((m for m in models if m.name == model_name), None)
    if not selected_model:
        return None, error_str(ValueError(f"Model '{model_name}' not found"))

    current_model = selected_model

    if progress:
        progress(0.1, desc=f"Using {model_name}")

    # Setup generator
    if seed <= 0:
        seed = torch.randint(0, 2**32-1, (1,)).item()
    generator = torch.Generator(device).manual_seed(seed)

    try:
        if mode == "img2img":
            result_image = img_to_img(
                current_model.path, prompt, negative_prompt, image, strength,
                guidance_scale, num_steps, width, height, generator,
                scale_method, scale_factor, progress
            )
        else:
            result_image = txt_to_img(
                current_model.path, prompt, negative_prompt, guidance_scale,
                num_steps, width, height, generator, scale_method, scale_factor,
                progress
            )

        if progress:
            progress(1.0, desc="Complete!")

        # Save result
        timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
        os.makedirs("imgs", exist_ok=True)
        filename = f"imgs/result-{timestamp}.png"
        result_image.save(filename)

        info = f"""### ✅ Generation Complete
- **Mode**: {mode}
- **Model**: {model_name}  
- **Resolution**: {result_image.size[0]}x{result_image.size[1]}
- **Scale**: {scale_factor}x ({scale_method})
- **Seed**: {seed}
- **Saved**: {filename}"""

        return result_image, info

    except Exception as e:
        print(f"Generation error: {e}")
        return None, error_str(e, "Generation Failed")

def txt_to_img(
    model_path: str, prompt: str, neg_prompt: str, guidance: float,
    steps: int, width: int, height: int, generator, scale: str,
    scale_factor: int, progress
) -> PIL.Image.Image:
    """Text-to-image generation"""

    global last_mode, pipe, current_model_path

    if progress:
        progress(0.2, desc="Loading pipeline...")

    # Load pipeline if needed
    if model_path != current_model_path or last_mode != "txt2img":
        current_model_path = model_path

        if is_colab or current_model == custom_model:
            pipe = StableDiffusionPipeline.from_pretrained(
                current_model_path,
                torch_dtype=dtype,
                scheduler=scheduler,
                safety_checker=None,
            )
        else:
            pipe = current_model.pipe_t2i

        if torch.cuda.is_available():
            pipe = pipe.to(device)
        last_mode = "txt2img"

    if progress:
        progress(0.4, desc="Generating image...")

    # Add model prefix
    full_prompt = f"{current_model.prefix}, {prompt}" if current_model.prefix else prompt

    result = pipe(
        full_prompt,
        negative_prompt=neg_prompt,
        num_inference_steps=int(steps),
        guidance_scale=guidance,
        width=width,
        height=height,
        generator=generator,
    ).images[0]

    if progress:
        progress(0.7, desc="Applying super resolution...")

    # Apply super resolution
    if scale_factor > 1:
        if scale == "RealESRGAN":
            fp32 = device == "cpu"
            result = realEsrgan(
                input_dir=result,
                suffix="",
                output_dir="imgs",
                fp32=fp32,
                outscale=scale_factor,
            )[0]
        else:  # Waifu2x
            result = magnifier.magnify(result, scale_factor=scale_factor)

    return result

def img_to_img(
    model_path: str, prompt: str, neg_prompt: str, img: PIL.Image.Image,
    strength: float, guidance: float, steps: int, width: int, height: int,
    generator, scale: str, scale_factor: int, progress
) -> PIL.Image.Image:
    """Image-to-image generation"""

    global last_mode, pipe, current_model_path

    if progress:
        progress(0.2, desc="Loading pipeline...")

    # Load pipeline if needed
    if model_path != current_model_path or last_mode != "img2img":
        current_model_path = model_path

        if is_colab or current_model == custom_model:
            pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
                current_model_path,
                torch_dtype=dtype,
                scheduler=scheduler,
                safety_checker=None,
            )
        else:
            pipe = current_model.pipe_i2i

        if torch.cuda.is_available():
            pipe = pipe.to(device)
        last_mode = "img2img"

    # Resize input image
    if progress:
        progress(0.3, desc="Processing input image...")

    ratio = min(height / img.height, width / img.width)
    img = img.resize((int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS)

    # Add model prefix
    full_prompt = f"{current_model.prefix}, {prompt}" if current_model.prefix else prompt

    if progress:
        progress(0.4, desc="Transforming image...")

    result = pipe(
        full_prompt,
        negative_prompt=neg_prompt,
        image=img,
        num_inference_steps=int(steps),
        strength=strength,
        guidance_scale=guidance,
        generator=generator,
    ).images[0]

    if progress:
        progress(0.7, desc="Applying super resolution...")

    # Apply super resolution
    if scale_factor > 1:
        if scale == "RealESRGAN":
            fp32 = device == "cpu"
            result = realEsrgan(
                input_dir=result,
                suffix="",
                output_dir="imgs",
                fp32=fp32,
                outscale=scale_factor,
            )[0]
        else:  # Waifu2x
            result = magnifier.magnify(result, scale_factor=scale_factor)

    return result

def load_example_images() -> List[str]:
    """Load example images for gallery"""
    example_images = []
    for f_img in findfile.find_cwd_files(".png", recursive=2):
        if "result-" in os.path.basename(f_img) or "random_examples" in f_img:
            example_images.append(f_img)
    return example_images[:12]  # Limit examples

# Custom CSS for styling
custom_css = """
.gradio-container {
    font-family: 'Segoe UI', system-ui, sans-serif;
    max-width: 1400px;
    margin: 0 auto;
}

.header-section {
    text-align: center;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    color: white;
    padding: 2rem;
    border-radius: 15px;
    margin-bottom: 2rem;
}

.info-card {
    background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
    color: white;
    padding: 1.5rem;
    border-radius: 10px;
    margin: 1rem 0;
}

.status-info {
    background: #e8f5e8;
    border-left: 4px solid #4CAF50;
    padding: 1rem;
    border-radius: 5px;
    margin: 1rem 0;
}

.generate-btn {
    background: linear-gradient(45deg, #FF6B6B, #4ECDC4) !important;
    border: none !important;
    border-radius: 25px !important;
    padding: 15px 30px !important;
    font-size: 16px !important;
    font-weight: bold !important;
    color: white !important;
    transition: all 0.3s ease !important;
}

.generate-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 10px 20px rgba(0,0,0,0.2) !important;
}
"""

def create_interface():
    """Create the Gradio interface"""

    with gr.Blocks(title=APP_TITLE, css=custom_css) as demo:

        # Header
        with gr.Row():
            gr.HTML(f"""
            <div class="header-section">
                <h1 style="font-size: 2.5rem; margin-bottom: 1rem;">{APP_TITLE}</h1>
                <p style="font-size: 1.2rem; margin-bottom: 1rem;">{APP_DESCRIPTION}</p>
                <div style="font-size: 1rem;">{get_system_info()}</div>
            </div>
            """)

        # Info sections
        with gr.Row():
            with gr.Column():
                gr.Markdown(INSTRUCTIONS, elem_classes=["info-card"])
            with gr.Column():
                gr.Markdown(CONTRIBUTION_INFO, elem_classes=["info-card"])

        # Copyright
        gr.Markdown(f"### 📄 {COPYRIGHT_INFO}", elem_classes=["status-info"])

        # Main interface
        with gr.Tabs():

            # Text-to-Image Tab
            with gr.TabItem("🎨 Text-to-Image"):
                with gr.Row():
                    with gr.Column(scale=1):

                        # Model selection
                        model_dropdown = gr.Dropdown(
                            choices=[m.name for m in models],
                            value=current_model.name,
                            label="🤖 Model Selection"
                        )

                        # Custom model path
                        custom_model_path = gr.Textbox(
                            label="🔗 Custom Model Path (HuggingFace)",
                            placeholder="username/model-name",
                            visible=custom_model is not None
                        )

                        # Prompts
                        prompt_txt2img = gr.Textbox(
                            label="✨ Prompt",
                            placeholder="Describe your desired image...",
                            lines=3
                        )

                        negative_prompt_txt2img = gr.Textbox(
                            label="🚫 Negative Prompt",
                            value=DEFAULT_NEGATIVE_PROMPT,
                            lines=2
                        )

                        # Parameters
                        with gr.Accordion("🎛️ Generation Parameters", open=False):
                            with gr.Row():
                                width_txt2img = gr.Slider(256, 1024, 512, step=64, label="Width")
                                height_txt2img = gr.Slider(256, 1024, 512, step=64, label="Height")

                            with gr.Row():
                                guidance_scale_txt2img = gr.Slider(1, 20, 7.5, step=0.5, label="Guidance Scale")
                                num_steps_txt2img = gr.Slider(10, 50, 20, label="Steps")

                            seed_txt2img = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)

                        # Super Resolution
                        with gr.Accordion("🔍 Super Resolution", open=True):
                            scale_method_txt2img = gr.Radio(
                                choices=["RealESRGAN", "Waifu2x"],
                                value="RealESRGAN",
                                label="Method"
                            )
                            scale_factor_txt2img = gr.Slider(1, 4, 2, step=1, label="Scale Factor")

                    with gr.Column(scale=1):
                        # Generate button
                        generate_btn_txt2img = gr.Button(
                            "🎨 Generate Image",
                            variant="primary",
                            elem_classes=["generate-btn"]
                        )

                        # Output
                        output_image_txt2img = gr.Image(label="Generated Image", type="pil")
                        output_info_txt2img = gr.Markdown("Ready to generate! 🚀")

            # Image-to-Image Tab
            with gr.TabItem("🖼️ Image-to-Image"):
                with gr.Row():
                    with gr.Column(scale=1):

                        # Input image
                        input_image_img2img = gr.Image(
                            label="📤 Input Image",
                            type="pil"
                        )

                        # Model selection
                        model_dropdown_img2img = gr.Dropdown(
                            choices=[m.name for m in models],
                            value=current_model.name,
                            label="🤖 Model Selection"
                        )

                        # Prompts
                        prompt_img2img = gr.Textbox(
                            label="✨ Transformation Prompt",
                            placeholder="How to transform the image...",
                            lines=3
                        )

                        negative_prompt_img2img = gr.Textbox(
                            label="🚫 Negative Prompt",
                            value=DEFAULT_NEGATIVE_PROMPT,
                            lines=2
                        )

                        # Parameters
                        with gr.Accordion("🎛️ Generation Parameters", open=False):
                            with gr.Row():
                                width_img2img = gr.Slider(256, 1024, 512, step=64, label="Width")
                                height_img2img = gr.Slider(256, 1024, 512, step=64, label="Height")

                            strength_img2img = gr.Slider(
                                0.1, 1.0, 0.75, step=0.05,
                                label="Strength (how much to change)"
                            )

                            with gr.Row():
                                guidance_scale_img2img = gr.Slider(1, 20, 7.5, step=0.5, label="Guidance")
                                num_steps_img2img = gr.Slider(10, 50, 20, label="Steps")

                            seed_img2img = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)

                        # Super Resolution
                        with gr.Accordion("🔍 Super Resolution", open=True):
                            scale_method_img2img = gr.Radio(
                                choices=["RealESRGAN", "Waifu2x"],
                                value="RealESRGAN",
                                label="Method"
                            )
                            scale_factor_img2img = gr.Slider(1, 4, 2, step=1, label="Scale Factor")

                    with gr.Column(scale=1):
                        # Generate button
                        generate_btn_img2img = gr.Button(
                            "🖼️ Transform Image",
                            variant="primary",
                            elem_classes=["generate-btn"]
                        )

                        # Output
                        output_image_img2img = gr.Image(label="Transformed Image", type="pil")
                        output_info_img2img = gr.Markdown("Upload an image to transform! 🖼️")

            # Gallery Tab
            with gr.TabItem("🖼️ Gallery"):
                gr.Markdown("### 🎨 Generated Images")

                with gr.Row():
                    refresh_gallery_btn = gr.Button("🔄 Refresh Gallery", variant="secondary")

                example_gallery = gr.Gallery(
                    value=load_example_images(),
                    label="Results Gallery",
                    show_label=False,
                    columns=4,
                    height="auto"
                )

        # Event handlers

        # Model changes
        model_dropdown.change(
            fn=on_model_change,
            inputs=[model_dropdown],
            outputs=[custom_model_path, prompt_txt2img]
        )

        # Sync models between tabs
        model_dropdown.change(
            fn=lambda x: gr.update(value=x),
            inputs=[model_dropdown],
            outputs=[model_dropdown_img2img]
        )

        model_dropdown_img2img.change(
            fn=lambda x: gr.update(value=x),
            inputs=[model_dropdown_img2img],
            outputs=[model_dropdown]
        )

        # Custom model path
        if custom_model:
            custom_model_path.change(
                fn=custom_model_changed,
                inputs=[custom_model_path],
                outputs=[output_info_txt2img]
            )

        # Generation events
        generate_btn_txt2img.click(
            fn=generate_image,
            inputs=[
                gr.State("txt2img"),
                model_dropdown,
                prompt_txt2img,
                negative_prompt_txt2img,
                width_txt2img,
                height_txt2img,
                guidance_scale_txt2img,
                num_steps_txt2img,
                seed_txt2img,
                gr.State(None),  # No input image for txt2img
                gr.State(0.75),  # Default strength
                scale_method_txt2img,
                scale_factor_txt2img
            ],
            outputs=[output_image_txt2img, output_info_txt2img]
        )

        generate_btn_img2img.click(
            fn=generate_image,
            inputs=[
                gr.State("img2img"),
                model_dropdown_img2img,
                prompt_img2img,
                negative_prompt_img2img,
                width_img2img,
                height_img2img,
                guidance_scale_img2img,
                num_steps_img2img,
                seed_img2img,
                input_image_img2img,
                strength_img2img,
                scale_method_img2img,
                scale_factor_img2img
            ],
            outputs=[output_image_img2img, output_info_img2img]
        )

        # Gallery refresh
        refresh_gallery_btn.click(
            fn=load_example_images,
            outputs=[example_gallery]
        )

    return demo

if __name__ == "__main__":
    # Initialize
    print(f"🚀 Starting {APP_TITLE}...")
    print(f"⏱️ Initialization time: {time.time() - start_time:.2f}s")
    print(f"🖥️ {get_system_info()}")

    # Ensure output directory
    os.makedirs("imgs", exist_ok=True)

    # Initialize models
    initialize_models()

    # Create and launch interface
    demo = create_interface()

    # Launch settings
    launch_kwargs = {
        "share": False,
        "server_name": "0.0.0.0",
        "server_port": 7860,
        "show_error": True,
    }

    if is_colab:
        launch_kwargs["share"] = True

    print("🌐 Launching WebUI...")
    demo.launch(**launch_kwargs)