import gradio as gr
from PIL import Image, ImageEnhance, ImageDraw
import os

# --- 1. Processing Logic (Updated with Preview) ---
def process_image(image, split_count, should_enhance, format_type):
    if image is None: return None, None # Return 2 nones (Preview & Gallery)
    
    # A. Enhancement
    if should_enhance:
        enhancer = ImageEnhance.Sharpness(image)
        image = enhancer.enhance(1.5)
        contrast = ImageEnhance.Contrast(image)
        image = contrast.enhance(1.1)
    
    img_width, img_height = image.size
    split_width = img_width // int(split_count)
    output_files = []
    
    # B. Creating the "Insta-Preview" (Stitching parts together)
    # Hum parts ko wapas jod kar dikhayenge taaki user ko 'Seamless' feel aaye
    preview_img = Image.new('RGB', (img_width, img_height))
    
    if not os.path.exists("outputs"): os.makedirs("outputs")
    
    for i in range(int(split_count)):
        left = i * split_width
        right = (i + 1) * split_width
        
        # Crop
        part = image.crop((left, 0, right, img_height))
        
        # Transparency Fix
        if part.mode in ('RGBA', 'LA') or (part.mode == 'P' and 'transparency' in part.info):
            part = part.convert('RGB')
        
        # Save for Gallery
        ext = "png" if format_type == "PNG (HQ)" else "jpg"
        save_path = f"outputs/ts_part_{i+1}.{ext}"
        part.save(save_path, quality=95)
        output_files.append(save_path)
        
        # Add to Preview (Paste side by side)
        preview_img.paste(part, (left, 0))
        
    return preview_img, output_files

# --- 2. New Premium Lighting CSS (SAME AS YOURS) ---
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;900&display=swap');

body, .gradio-container {
    font-family: 'Inter', sans-serif !important;
    background: #0a0a0a !important; /* Deep Dark Background */
    color: #ffffff !important;
}

/* The Main Card with Lighting Effect */
#main_card {
    background: rgba(30, 30, 40, 0.8);
    backdrop-filter: blur(30px);
    -webkit-backdrop-filter: blur(30px);
    border: 1px solid rgba(255, 255, 255, 0.15);
    border-radius: 30px;
    padding: 50px;
    box-shadow: 0 0 50px rgba(100, 200, 255, 0.15), inset 0 0 20px rgba(255, 255, 255, 0.05);
    max-width: 1100px;
    margin: 40px auto;
}

/* New TS Logo Style */
#logo_text h1 {
    color: #ffffff;
    font-size: 5rem !important;
    font-weight: 900;
    text-align: center;
    margin-bottom: 0px;
    letter-spacing: -2px;
    text-shadow: 0 0 25px rgba(150, 220, 255, 0.8);
}

#subtitle_text h3 {
    text-align: center;
    color: #bbbbbb;
    font-size: 1.3rem;
    font-weight: 400;
    margin-top: 15px;
    margin-bottom: 50px;
    letter-spacing: 1px;
}

/* Input Controls Styling */
.block, .svelte-12cmxck, .form {
    background: rgba(255,255,255,0.03) !important;
    border: 1px solid rgba(255,255,255,0.1) !important;
    border-radius: 20px !important;
    padding: 20px !important;
}

.gradio-row { gap: 25px !important; flex-wrap: wrap; }
label span { color: #fff !important; font-weight: 600; font-size: 1rem;}

/* The Button - Glowing */
.primary-btn {
    background: linear-gradient(135deg, #ffffff, #d0d0d0) !important;
    color: #000 !important;
    border: none !important;
    font-size: 1.3rem !important;
    font-weight: 900 !important;
    padding: 20px !important;
    border-radius: 20px !important;
    box-shadow: 0 0 30px rgba(255, 255, 255, 0.3);
    transition: all 0.3s ease;
}
.primary-btn:hover {
    transform: translateY(-3px);
    box-shadow: 0 0 50px rgba(255, 255, 255, 0.5);
}
"""

# --- 3. Spacious Interface Structure ---
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as app:
    
    with gr.Column(elem_id="main_card"):
        gr.Markdown("# TS", elem_id="logo_text")
        gr.Markdown("### PREMIUM CAROUSEL STUDIO", elem_id="subtitle_text")
        
        with gr.Row():
            # Left Side (Inputs)
            with gr.Column(scale=1):
                inp_img = gr.Image(type="pil", label="1. Drop Your Wide Image", sources=["upload", "clipboard"], height=300)
                
                with gr.Column():
                    slider = gr.Slider(2, 5, value=3, step=1, label="2. How many parts?")
                    
                    with gr.Row():
                        enhance_opt = gr.Checkbox(label="Magic Enhance", value=True)
                        fmt_opt = gr.Radio(["JPG", "PNG (HQ)"], label="Format", value="JPG")
                
                btn = gr.Button("⚡ GENERATE & PREVIEW", variant="primary", elem_classes=["primary-btn"])

            # Right Side (Output - Now with Preview)
            with gr.Column(scale=1.5):
                # Ye hai wo naya FEATURE:
                preview_out = gr.Image(label="📱 Insta-Preview (Check Seamless Look)", interactive=False)
                out_gal = gr.Gallery(label="⬇️ Download Parts Here", columns=3, height="auto", object_fit="contain")

    # Connect both outputs (Preview + Gallery)
    btn.click(process_image, inputs=[inp_img, slider, enhance_opt, fmt_opt], outputs=[preview_out, out_gal])

app.launch()