import gradio as gr
from gradio_client import Client, handle_file
from PIL import Image

# Connect to the verified IC-Light engine
client = Client("lllyasviel/IC-Light")

def get_smart_resolution(width, height):
    """
    Auto-detects Landscape/Portrait to prevent stretching.
    """
    aspect_ratio = width / height
    if aspect_ratio > 1.2:  # Landscape (Like your new photo)
        return 768, 512
    elif aspect_ratio < 0.8: # Portrait
        return 512, 768
    else: # Square
        return 640, 640

def dynamic_relight(image_path, prompt, lighting_choice, s1, s2, s3, s4):
    if image_path is None:
        return None

    # 1. SMART RESOLUTION
    img = Image.open(image_path)
    orig_w, orig_h = img.size
    target_w, target_h = get_smart_resolution(orig_w, orig_h)
    
    img.thumbnail((target_w, target_h)) 
    safe_path = "robust_input.png"
    img.save(safe_path)

    # 2. CALIBRATED SLIDERS (Visual Enhancement ONLY)
    # We removed "Cinematic" from the hidden logic to prevent darkness.
    # Now sliders only add QUALITY (Sharpness, Texture), not Darkness.
    keywords = ["ultra detailed", "sharp texture", "realistic skin", "perfect focus"]
    slider_vals = [s1, s2, s3, s4]
    
    user_adjustments = ""
    for i, word in enumerate(keywords):
        if slider_vals[i] > 0:
            weight = 0.5 + (slider_vals[i] / 100.0) * 0.7  # Gentle range (0.5 to 1.2)
            user_adjustments += f", ({word}:{weight:.1f})"
    
    # 3. THE "DAYLIGHT PROTECTION" GUARD
    # This acts as a firewall. It blocks the model from making the image dark.
    # It forces 'global illumination' so the background stays visible.
    lighting_guard = ", natural lighting, ambient light, global illumination, balanced exposure, detailed background, raw photo"
    
    # We combine your prompt + The Guard + The Sliders
    final_prompt = prompt + lighting_guard + user_adjustments

    try:
        result = client.predict(
            input_fg=handle_file(safe_path),   
            prompt=final_prompt,               
            image_width=target_w,              
            image_height=target_h,             
            num_samples=1,                     
            seed=12345,                        
            steps=30,                          
            a_prompt="best quality, masterpiece, 8k uhd", # Pure Quality boosters           
            n_prompt="lowres, bad anatomy, dark, moody, shadow, silhouette, black background", # ANTI-DARKNESS Negative Prompts
            cfg=1.8,                           # Slightly lower CFG prevents "hallucinating" new backgrounds
            highres_scale=1.5,                 
            highres_denoise=0.5,               
            lowres_denoise=0.9,                
            bg_source=lighting_choice,         
            api_name="/process_relight"        
        )
        return result[1][0]['image']
        
    except Exception as e:
        raise gr.Error(f"System Error: {str(e)}")

# UI Setup
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# ☀️ Robust Daylight Relighter")
    gr.Markdown("Guarantees bright, high-quality results for any input image.")
    
    with gr.Row():
        with gr.Column():
            img = gr.Image(type="filepath", label="Input Image")
            
            # DEFAULT PROMPT is now generic and safe
            txt = gr.Textbox(label="Prompt (Optional)", value="beautiful woman, detailed face")
            
            dirs = gr.Radio(["Left Light", "Right Light", "Top Light", "Bottom Light"], label="Light Direction", value="Top Light")
            
            gr.Markdown("### Enhancement (Safe Range)")
            # Renamed 'Cinematic' to 'Contrast' to be accurate
            s1 = gr.Slider(0, 100, value=20, label="Contrast Boost") 
            s2 = gr.Slider(0, 100, value=80, label="Detail Level")
            s3 = gr.Slider(0, 100, value=70, label="Texture Sharpness")
            s4 = gr.Slider(0, 100, value=60, label="Focus Depth")
            
            btn = gr.Button("Execute Robust Relight", variant="primary")
        out = gr.Image(label="Enhanced Output")

    btn.click(dynamic_relight, [img, txt, dirs, s1, s2, s3, s4], out)

demo.launch(show_error=True)