import gradio as gr
import torch
from PIL import Image
import numpy as np
from transformers import pipeline
import requests
from io import BytesIO
import os
from huggingface_hub import login
import warnings
warnings.filterwarnings("ignore")

class PhotoUpscaler:
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.current_model = None
        self.upscaler = None
        self.load_default_model()
    
    def load_default_model(self):
        """Load default upscaling model with Real-ESRGAN priority"""
        # Priority list of models to try
        priority_models = [
            "ai-forever/Real-ESRGAN",
            "sberbank-ai/Real-ESRGAN", 
            "caidas/swin2SR-realworld-sr-x4-64-bsrgan-psnr",
            "microsoft/swin2SR-compressed-sr-x2-48"
        ]
        
        for model_name in priority_models:
            try:
                self.current_model = model_name
                if "Real-ESRGAN" in model_name:
                    # Special handling for Real-ESRGAN models
                    try:
                        from diffusers import StableDiffusionUpscalePipeline
                        self.upscaler = StableDiffusionUpscalePipeline.from_pretrained(
                            model_name, 
                            torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
                        ).to(self.device)
                        return f"✅ Real-ESRGAN model načten: {self.current_model}"
                    except:
                        # Fallback to regular pipeline
                        self.upscaler = pipeline(
                            "image-to-image",
                            model=model_name,
                            device=0 if self.device == "cuda" else -1
                        )
                        return f"✅ Model načten: {self.current_model}"
                else:
                    # Regular Swin2SR models
                    self.upscaler = pipeline(
                        "image-to-image",
                        model=model_name,
                        device=0 if self.device == "cuda" else -1
                    )
                    return f"✅ Model načten: {self.current_model}"
            except Exception as e:
                print(f"Nepodařilo se načíst {model_name}: {e}")
                continue
        
        return f"❌ Nepodařilo se načíst žádný model"
    
    def upscale_image(self, image, scale_factor=2, model_choice="default"):
        """Upscale image using selected model"""
        if image is None:
            return None, "❌ Žádný obrázek nebyl nahrán"
        
        try:
            # Convert to PIL if needed
            if isinstance(image, np.ndarray):
                image = Image.fromarray(image)
            
            # Resize for processing if image is too large
            max_size = 1024
            if max(image.size) > max_size:
                ratio = max_size / max(image.size)
                new_size = tuple(int(dim * ratio) for dim in image.size)
                image = image.resize(new_size, Image.Resampling.LANCZOS)
            
            # Change model if requested
            if model_choice != "default" and model_choice != self.current_model:
                self.load_model(model_choice)
            
            # Perform upscaling based on model type
            if self.upscaler:
                try:
                    if "Real-ESRGAN" in self.current_model:
                        # Special handling for Real-ESRGAN models
                        if hasattr(self.upscaler, '__call__'):
                            # Diffusers pipeline
                            prompt = "high quality, detailed, sharp"
                            upscaled = self.upscaler(
                                prompt=prompt, 
                                image=image,
                                num_inference_steps=20,
                                guidance_scale=0
                            ).images[0]
                        else:
                            # Regular pipeline fallback
                            upscaled = self.upscaler(image)
                    elif "stable-diffusion" in self.current_model.lower():
                        # Stable Diffusion upscaler
                        from diffusers import StableDiffusionUpscalePipeline
                        prompt = "high quality, detailed, sharp, realistic"
                        upscaled = self.upscaler(
                            prompt=prompt,
                            image=image,
                            num_inference_steps=20
                        ).images[0]
                    else:
                        # Standard Swin2SR and other models
                        upscaled = self.upscaler(image)
                        if isinstance(upscaled, list):
                            upscaled = upscaled[0]
                        if hasattr(upscaled, 'images'):
                            upscaled = upscaled.images[0]
                        elif isinstance(upscaled, dict) and 'image' in upscaled:
                            upscaled = upscaled['image']
                    
                    return upscaled, f"✅ Obrázek zvětšen pomocí {MODEL_DESCRIPTIONS.get(self.current_model, self.current_model)}"
                    
                except Exception as model_error:
                    print(f"Model error: {model_error}")
                    # Fallback to simple upscaling
                    new_size = tuple(int(dim * scale_factor) for dim in image.size)
                    upscaled = image.resize(new_size, Image.Resampling.LANCZOS)
                    return upscaled, f"✅ Obrázek zvětšen pomocí klasického algoritmu (model selhání)"
            else:
                # Simple fallback upscaling
                new_size = tuple(int(dim * scale_factor) for dim in image.size)
                upscaled = image.resize(new_size, Image.Resampling.LANCZOS)
                return upscaled, f"✅ Obrázek zvětšen pomocí klasického algoritmu (fallback)"
                
        except Exception as e:
            return None, f"❌ Chyba při zpracování: {str(e)}"
    
    def load_model(self, model_name):
        """Load specific model with enhanced support for different model types"""
        try:
            self.current_model = model_name
            
            if "Real-ESRGAN" in model_name:
                # Try Real-ESRGAN specific loading
                try:
                    from diffusers import StableDiffusionUpscalePipeline
                    self.upscaler = StableDiffusionUpscalePipeline.from_pretrained(
                        model_name,
                        torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
                    ).to(self.device)
                    return f"✅ Real-ESRGAN model načten: {MODEL_DESCRIPTIONS.get(model_name, model_name)}"
                except:
                    # Fallback to regular pipeline
                    self.upscaler = pipeline(
                        "image-to-image",
                        model=model_name,
                        device=0 if self.device == "cuda" else -1
                    )
                    return f"✅ Model načten (fallback): {MODEL_DESCRIPTIONS.get(model_name, model_name)}"
            
            elif "stable-diffusion" in model_name.lower():
                # Stable Diffusion upscaler
                from diffusers import StableDiffusionUpscalePipeline
                self.upscaler = StableDiffusionUpscalePipeline.from_pretrained(
                    model_name,
                    torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
                ).to(self.device)
                return f"✅ SD Upscaler načten: {MODEL_DESCRIPTIONS.get(model_name, model_name)}"
            
            elif "ldm" in model_name.lower():
                # LDM models
                from diffusers import LDMSuperResolutionPipeline
                self.upscaler = LDMSuperResolutionPipeline.from_pretrained(
                    model_name,
                    torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
                ).to(self.device)
                return f"✅ LDM model načten: {MODEL_DESCRIPTIONS.get(model_name, model_name)}"
            
            else:
                # Standard pipeline for Swin2SR and similar models
                self.upscaler = pipeline(
                    "image-to-image",
                    model=model_name,
                    device=0 if self.device == "cuda" else -1
                )
                return f"✅ Model načten: {MODEL_DESCRIPTIONS.get(model_name, model_name)}"
                
        except Exception as e:
            return f"❌ Chyba při načítání modelu {MODEL_DESCRIPTIONS.get(model_name, model_name)}: {str(e)}"

# Initialize upscaler
upscaler = PhotoUpscaler()

# Available models for upscaling
UPSCALING_MODELS = [
    "default",
    # Real-ESRGAN models (nejlepší pro realistické fotografie)
    "ai-forever/Real-ESRGAN",
    "sberbank-ai/Real-ESRGAN",
    "philz1337x/clarity-upscaler",
    
    # BSRGAN models (vynikající pro reálné obrázky)
    "caidas/swin2SR-realworld-sr-x4-64-bsrgan-psnr",
    "caidas/swin2SR-realworld-sr-x2-64-bsrgan-psnr",
    
    # SwinIR models (state-of-the-art)
    "caidas/swinIR-M-real-sr-x4-64-bsrgan-psnr",
    "caidas/swinIR-L-real-sr-x4-64-bsrgan-psnr",
    
    # Microsoft Swin2SR (optimalizované)
    "microsoft/swin2SR-compressed-sr-x2-48",
    "microsoft/swin2SR-compressed-sr-x4-48",
    "microsoft/swin2SR-classical-sr-x2-64",
    "microsoft/swin2SR-classical-sr-x4-64",
    "microsoft/swin2SR-realworld-sr-x4-64-bsrgan-psnr",
    
    # Další pokročilé modely
    "Kolors/Kolors-IP-Adapter-FaceID-Plus",
    "stabilityai/stable-diffusion-x4-upscaler",
    "CompVis/ldm-super-resolution-4x-openimages"
]

# Model descriptions for better user experience
MODEL_DESCRIPTIONS = {
    "default": "🎯 Výchozí model - rychlý a spolehlivý",
    "ai-forever/Real-ESRGAN": "🏆 Real-ESRGAN - nejlepší pro fotografie",
    "sberbank-ai/Real-ESRGAN": "⭐ Real-ESRGAN Sberbank - vylepšená verze",
    "philz1337x/clarity-upscaler": "✨ Clarity Upscaler - ultra ostrý",
    "caidas/swin2SR-realworld-sr-x4-64-bsrgan-psnr": "🌟 BSRGAN 4x - premium kvalita",
    "caidas/swin2SR-realworld-sr-x2-64-bsrgan-psnr": "🌟 BSRGAN 2x - rychlejší",
    "caidas/swinIR-M-real-sr-x4-64-bsrgan-psnr": "🚀 SwinIR Medium - vyváženost",
    "caidas/swinIR-L-real-sr-x4-64-bsrgan-psnr": "🔥 SwinIR Large - maximální kvalita",
    "microsoft/swin2SR-compressed-sr-x2-48": "⚡ Komprimovaný 2x - rychlý",
    "microsoft/swin2SR-compressed-sr-x4-48": "⚡ Komprimovaný 4x - rychlý",
    "microsoft/swin2SR-classical-sr-x2-64": "🎨 Klasický 2x - digitální obrázky",
    "microsoft/swin2SR-classical-sr-x4-64": "🎨 Klasický 4x - digitální obrázky",
    "stabilityai/stable-diffusion-x4-upscaler": "🎭 SD Upscaler - kreativní vylepšení",
    "CompVis/ldm-super-resolution-4x-openimages": "🧠 LDM - generativní upscaling"
}

def process_upscaling(image, scale_factor, model_choice, hf_token):
    """Main processing function"""
    # Login to HuggingFace if token provided
    if hf_token and hf_token.strip():
        try:
            login(hf_token)
            status_msg = "🔐 Přihlášen k Hugging Face | "
        except:
            status_msg = "⚠️ Problém s HF tokenem | "
    else:
        status_msg = "ℹ️ Používám veřejné modely | "
    
    # Perform upscaling
    result_image, process_msg = upscaler.upscale_image(image, scale_factor, model_choice)
    
    return result_image, status_msg + process_msg

def get_model_info():
    """Get current model information"""
    device_info = f"Zařízení: {upscaler.device.upper()}"
    model_info = f"Aktuální model: {upscaler.current_model}"
    return f"ℹ️ {device_info} | {model_info}"

# Create Gradio interface
with gr.Blocks(
    title="🚀 Photo Upscaler - Hugging Face",
    theme=gr.themes.Soft(),
    css="""
    .gradio-container {
        max-width: 1200px !important;
        margin: auto !important;
    }
    .title {
        text-align: center;
        color: #ff6b35;
        margin-bottom: 20px;
    }
    """
) as demo:
    
    gr.HTML("""
    <div class="title">
        <h1>🚀 Photo Upscaler s Hugging Face</h1>
        <p>Zvětšujte své fotografie pomocí pokročilých AI modelů</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Vstup")
            input_image = gr.Image(
                label="Nahrajte fotografii",
                type="pil",
                format="png"
            )
            
            scale_factor = gr.Slider(
                minimum=1.5,
                maximum=4.0,
                value=2.0,
                step=0.5,
                label="Faktor zvětšení",
                info="Kolikrát zvětšit obrázek"
            )
            
            model_choice = gr.Dropdown(
                choices=[(MODEL_DESCRIPTIONS.get(model, model), model) for model in UPSCALING_MODELS],
                value="default",
                label="Vyberte model",
                info="Různé modely pro různé typy obrázků - Real-ESRGAN nejlepší pro fotografie"
            )
            
            hf_token = gr.Textbox(
                label="Hugging Face Token (volitelné)",
                placeholder="hf_xxxxxxxxxxxxx",
                type="password",
                info="Pro přístup k privátním modelům"
            )
            
            upscale_btn = gr.Button(
                "🔍 Zvětšit obrázek",
                variant="primary",
                size="lg"
            )
            
        with gr.Column(scale=1):
            gr.Markdown("### 📥 Výstup")
            output_image = gr.Image(
                label="Zvětšený obrázek",
                type="pil"
            )
            
            status_text = gr.Textbox(
                label="Status",
                interactive=False,
                max_lines=3
            )
            
            info_btn = gr.Button("ℹ️ Info o modelu")
    
    # Event handlers
    upscale_btn.click(
        fn=process_upscaling,
        inputs=[input_image, scale_factor, model_choice, hf_token],
        outputs=[output_image, status_text]
    )
    
    info_btn.click(
        fn=get_model_info,
        outputs=status_text
    )
    
    # Examples and tips
    gr.Markdown("### 📋 Tipy pro nejlepší výsledky")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("""
            **🏆 Nejlepší modely pro fotografie:**
            - **Real-ESRGAN**: Nejkvalitnější pro reálné fotky
            - **BSRGAN**: Vynikající detail a ostrost
            - **SwinIR Large**: Maximální kvalita, pomalejší
            - **Clarity Upscaler**: Ultra ostrý výsledek
            """)
        
        with gr.Column():
            gr.Markdown("""
            **⚡ Rychlé modely:**
            - **Komprimované modely**: Rychlé zpracování
            - **2x modely**: Rychlejší než 4x verze
            - **Classical modely**: Pro digitální obrázky
            """)
    
    gr.Markdown("""
    ### 💡 Doporučení podle typu obrázku:
    - **Portréty**: Real-ESRGAN nebo SwinIR Large
    - **Krajiny**: BSRGAN nebo Clarity Upscaler  
    - **Staré fotky**: Real-ESRGAN s noise reduction
    - **Digitální art**: Classical nebo Stable Diffusion Upscaler
    - **Dokumenty**: SwinIR Medium pro čitelnost
    
    ### ⚙️ Optimalizace výkonu:
    - **GPU**: Automaticky detekováno pro rychlejší zpracování
    - **Velikost**: 256-512px pro nejlepší poměr rychlost/kvalita
    - **Formát**: PNG zachovává nejvyšší kvalitu
    - **HF Token**: Pro přístup k nejnovějším modelům
    """)

if __name__ == "__main__":
    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )