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"""
        try:
            # Using a Real-ESRGAN style model from Hugging Face
            self.current_model = "caidas/swin2SR-realworld-sr-x4-64-bsrgan-psnr"
            self.upscaler = pipeline(
                "image-to-image",
                model=self.current_model,
                device=0 if self.device == "cuda" else -1
            )
            return f"✅ Model načten: {self.current_model}"
        except Exception as e:
            # Fallback to a simpler approach
            self.current_model = "microsoft/swin2SR-compressed-sr-x2-48"
            try:
                self.upscaler = pipeline(
                    "image-to-image", 
                    model=self.current_model,
                    device=0 if self.device == "cuda" else -1
                )
                return f"✅ Fallback model načten: {self.current_model}"
            except:
                return f"❌ Chyba při načítání modelů: {str(e)}"
    
    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
            if self.upscaler:
                # For pipeline-based upscaling
                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í {self.current_model}"
            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"""
        try:
            self.current_model = model_name
            self.upscaler = pipeline(
                "image-to-image",
                model=model_name,
                device=0 if self.device == "cuda" else -1
            )
            return f"✅ Model změněn na: {model_name}"
        except Exception as e:
            return f"❌ Chyba při načítání modelu {model_name}: {str(e)}"

# Initialize upscaler
upscaler = PhotoUpscaler()

# Available models for upscaling
UPSCALING_MODELS = [
    "default",
    "microsoft/swin2SR-compressed-sr-x2-48",
    "microsoft/swin2SR-compressed-sr-x4-48", 
    "caidas/swin2SR-realworld-sr-x4-64-bsrgan-psnr",
    "microsoft/swin2SR-classical-sr-x2-64",
    "microsoft/swin2SR-classical-sr-x4-64"
]

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=UPSCALING_MODELS,
                value="default",
                label="Vyberte model",
                info="Různé modely pro různé typy obrázků"
            )
            
            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
    gr.Markdown("### 📋 Tipy pro použití")
    gr.Markdown("""
    - **Nejlepší výsledky**: Používejte obrázky s rozlišením 256x256 až 512x512 pixelů
    - **Modely**: Různé modely jsou optimalizované pro různé typy obrázků
    - **Real-world modely**: Nejlepší pro fotografie z reálného světa
    - **Classical modely**: Vhodné pro umělé nebo digitální obrázky
    - **HF Token**: Zadejte 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
    )