import gradio as gr
import os
import tempfile
import zipfile
import shutil
from test_api import run_adain

def process_adain_api(files, style_strength=1.0, dataset_size=100):
    '''API endpoint for AdaIN processing.'''
    if not files:
        return None, "No files uploaded"
    
    # Create temporary directories
    temp_dir = tempfile.mkdtemp()
    input_dir = os.path.join(temp_dir, "input")
    output_dir = os.path.join(temp_dir, "output")
    
    os.makedirs(input_dir, exist_ok=True)
    os.makedirs(output_dir, exist_ok=True)
    
    # Save uploaded files
    for file in files:
        if file is not None:
            shutil.copy(file.name, input_dir)
    
    try:
        style_dataset_pth = "tidalove/paleo-real"
        
        # Run AdaIN
        run_adain(input_dir, style_dataset_pth, output_dir, style_strength, dataset_size)
        
        # Create zip file
        zip_path = os.path.join(temp_dir, "style_transfer_results.zip")
        with zipfile.ZipFile(zip_path, 'w') as zipf:
            for file in os.listdir(output_dir):
                if file.lower().endswith(('.jpg', '.jpeg', '.png')):
                    zipf.write(os.path.join(output_dir, file), file)
        
        return [os.path.join(output_dir, basename) for basename in os.listdir(output_dir)], zip_path, f"Style transfer completed with strength {style_strength}"
        
    except Exception as e:
        return None, f"Error: {str(e)}"

# Create interface with API endpoint
with gr.Blocks() as adain_demo:
    gr.Markdown("# AdaIN Style Transfer Service")
    
    with gr.Row():
        with gr.Column():
            files_input = gr.File(label="Upload Images", file_count="multiple", file_types=["image"])
            strength_input = gr.Slider(0.0, 2.0, 1.0, step=0.1, label="Style strength")
            size_input = gr.Slider(1, 1000, 100, step=10, label="Generated dataset size")
            process_btn = gr.Button("Process", variant="primary")
            
        with gr.Column():
            gallery_output = gr.Gallery(label="Gallery", columns=3, rows=2, object_fit="contain")
            download_output = gr.File(label="Download Results")
            status_output = gr.Textbox(label="Status", interactive=False)
    
    process_btn.click(
        fn=process_adain_api,
        inputs=[files_input, strength_input, size_input],
        outputs=[gallery_output, download_output, status_output],
        api_name="adain_process"  # This creates the API endpoint
    )

adain_demo.launch(show_error=True)