import gradio as gr
import os
from pathlib import Path

# Add error handling for Hugging Face environment
try:
    from sky_masking import SkyMaskingPipeline
    from sky_replacement import StateOfTheArtSkyReplacer
except ImportError as e:
    print(f"Import error: {e}")
    # Handle missing dependencies gracefully

class SkyReplacementApp:
    def __init__(self):
        # Initialize with error handling
        try:
            self.sky_masker = SkyMaskingPipeline()
            self.sky_replacer = StateOfTheArtSkyReplacer()
            print("✅ App initialized successfully!")
        except Exception as e:
            print(f"❌ Error initializing app: {e}")

    def process_image(self, input_image):
        try:
            # Your existing processing logic
            sky_mask = self.sky_masker.generate_mask(input_image)
            result_image = self.sky_replacer.replace_sky(input_image, sky_mask)
            return result_image
        except Exception as e:
            print(f"❌ Error processing image: {str(e)}")
            return input_image

def create_interface():
    app = SkyReplacementApp()
    
    interface = gr.Interface(
        fn=app.process_image,
        inputs=gr.Image(label="Upload Image", type="pil"),
        outputs=gr.Image(label="Sky Replaced Result", type="pil"),
        title="🌤️ AI Sky Replacement - 2025 State-of-the-Art",
        description="Upload an image to replace its sky with premium-quality alternatives using advanced edge refinement.",
        examples=[
            # Add example images if available
        ],
        theme="default"
    )
    
    return interface

if __name__ == "__main__":
    demo = create_interface()
    # For Hugging Face Spaces
    demo.launch(share=False)