import os
import cv2
import gradio as gr
import numpy as np
import random
import requests
import torch
from torchvision import transforms
from PIL import Image
import time

# Simulated style transfer function - in real app, replace with actual ML model
def apply_style_transfer(content_img, style_img, style_strength, randomize_strength):
    if content_img is None or style_img is None:
        return None, None, "Empty image"
    
    if randomize_strength:
        style_strength = random.uniform(0.1, 1.0)
    
    # Simulate processing time
    time.sleep(3)
    
    # In a real app, this would be your actual style transfer logic
    # For demo, we'll just blend the images
    content_img = cv2.resize(content_img, (512, 512))
    style_img = cv2.resize(style_img, (512, 512))
    result = cv2.addWeighted(content_img, 1-style_strength, style_img, style_strength, 0)
    
    return result, style_strength, "Success"

# Constants
MAX_STYLE_STRENGTH = 1.0
MIN_STYLE_STRENGTH = 0.1

example_path = os.path.join(os.path.dirname(__file__), 'assets')

# Assume we have these directories with example images
content_list = os.listdir(os.path.join(example_path, "content"))
content_list_path = [os.path.join(example_path, "content", img) for img in content_list]

style_list = os.listdir(os.path.join(example_path, "styles"))
style_list_path = [os.path.join(example_path, "styles", style) for style in style_list]

css = """
#col-content {
    margin: 0 auto;
    max-width: 430px;
}
#col-style {
    margin: 0 auto;
    max-width: 430px;
}
#col-result {
    margin: 0 auto;
    max-width: 430px;
}
#gallery {
    margin: 0 auto;
    max-width: 1100px;
}
#transfer-button {
    color: purple;
    font-size: 18px;
}
"""

with gr.Blocks(css=css) as ArtStyleTransfer:
    gr.Markdown("# 🎨 AI Art Style Transfer")
    gr.Markdown("Transform your photos into artistic masterpieces!")
    
    with gr.Row():
        with gr.Column(elem_id="col-content"):
            gr.Markdown("### 1. Choose Your Content Image")
            content_img = gr.Image(label="Photo to transform", sources='upload', type="numpy")
            content_examples = gr.Examples(
                inputs=content_img,
                examples=content_list_path,
                examples_per_page=8
            )
        
        with gr.Column(elem_id="col-style"):
            gr.Markdown("### 2. Pick an Artistic Style")
            style_img = gr.Image(label="Style reference", sources='upload', type="numpy")
            style_examples = gr.Examples(
                inputs=style_img,
                examples=style_list_path,
                examples_per_page=8
            )
        
        with gr.Column(elem_id="col-result"):
            gr.Markdown("### 3. Generate Artwork")
            result_img = gr.Image(label="Stylized Result")
            with gr.Row():
                style_strength = gr.Slider(
                    label="Style Strength",
                    minimum=MIN_STYLE_STRENGTH,
                    maximum=MAX_STYLE_STRENGTH,
                    step=0.1,
                    value=0.5,
                )
                randomize_strength = gr.Checkbox(label="Randomize Strength")
            
            with gr.Row():
                final_strength = gr.Number(label="Applied Style Strength")
                status = gr.Text(label="Status")
            
            transfer_button = gr.Button("🎨 Transform!", elem_id="transfer-button")
    
    transfer_button.click(
        fn=apply_style_transfer,
        inputs=[content_img, style_img, style_strength, randomize_strength],
        outputs=[result_img, final_strength, status],
        api_name=False,
        concurrency_limit=20
    )
    
    with gr.Column(elem_id="gallery"):
        gr.Markdown("## ✨ Style Transfer Gallery")
        show_case = gr.Examples(
            examples=[
                ["assets/examples/photo1.jpg", "assets/examples/style1.jpg", "assets/examples/result1.jpg"],
                ["assets/examples/photo2.jpg", "assets/examples/style2.jpg", "assets/examples/result2.jpg"],
            ],
            inputs=[content_img, style_img, result_img],
            label=None
        )

ArtStyleTransfer.queue(api_open=False).launch(show_api=False)