import gradio as gr
import torch
from PIL import Image
import numpy as np
from transformers import pipeline
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 = "microsoft/swin2SR-compressed-sr-x2-48"
        self.upscaler = None
        self.load_default_model()
    
    def load_default_model(self):
        """Load default upscaling model optimized for HF Spaces"""
        try:
            self.upscaler = pipeline(
                "image-to-image",
                model=self.current_model,
                device=0 if self.device == "cuda" else -1
            )
            return f"✅ Model loaded: {self.current_model}"
        except Exception as e:
            return f"❌ Error loading model: {str(e)}"
    
    def upscale_image(self, image, scale_factor=2):
        """Upscale image with simplified processing for HF Spaces"""
        if image is None:
            return None, "❌ No image uploaded"
        
        try:
            # Convert to PIL if needed
            if isinstance(image, np.ndarray):
                image = Image.fromarray(image)
            
            # Limit size for HF Spaces (smaller memory limits)
            max_size = 512
            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)
            
            # Perform upscaling
            if self.upscaler:
                try:
                    upscaled = self.upscaler(image)
                    if isinstance(upscaled, list):
                        upscaled = upscaled[0]
                    return upscaled, f"✅ Image upscaled using {self.current_model}"
                except:
                    # Fallback to simple upscaling if model fails
                    new_size = tuple(int(dim * scale_factor) for dim in image.size)
                    upscaled = image.resize(new_size, Image.Resampling.LANCZOS)
                    return upscaled, "✅ Image upscaled using fallback method"
            else:
                new_size = tuple(int(dim * scale_factor) for dim in image.size)
                upscaled = image.resize(new_size, Image.Resampling.LANCZOS)
                return upscaled, "✅ Image upscaled using classic algorithm"
                
        except Exception as e:
            return None, f"❌ Processing error: {str(e)}"

# Initialize upscaler
upscaler = PhotoUpscaler()

def process_upscaling(image, scale_factor):
    """Main processing function for HF Spaces"""
    result_image, process_msg = upscaler.upscale_image(image, scale_factor)
    return result_image, process_msg

# Create Gradio interface optimized for HF Spaces
with gr.Blocks(
    title="🚀 Photo Upscaler",
    theme=gr.themes.Default(),
    css="""
    .gradio-container {
        max-width: 900px !important;
        margin: auto !important;
    }
    """
) as demo:
    
    gr.HTML("""
    <div style="text-align: center; margin-bottom: 20px;">
        <h1>🚀 Photo Upscaler</h1>
        <p>Enhance your photos with AI-powered upscaling</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### 📤 Input")
            input_image = gr.Image(
                label="Upload Photo",
                type="pil"
            )
            
            scale_factor = gr.Slider(
                minimum=1.5,
                maximum=3.0,
                value=2.0,
                step=0.5,
                label="Scale Factor"
            )
            
            upscale_btn = gr.Button(
                "🔍 Upscale Image",
                variant="primary"
            )
            
        with gr.Column():
            gr.Markdown("### 📥 Output")
            output_image = gr.Image(
                label="Upscaled Image"
            )
            
            status_text = gr.Textbox(
                label="Status",
                interactive=False
            )
    
    # Event handler
    upscale_btn.click(
        fn=process_upscaling,
        inputs=[input_image, scale_factor],
        outputs=[output_image, status_text]
    )
    
    # Tips
    gr.Markdown("""
    ### 💡 Tips for best results:
    - Use images between 256x256 and 512x512 pixels
    - PNG format provides best quality
    - Lower scale factors work faster
    """)

# Launch for HF Spaces
if __name__ == "__main__":
    demo.launch()