import gradio as gr
from transformers import pipeline
from PIL import Image
import os
from core_processor import SculptorEngine

# --- 1. Load Models (GPU Start hote hi load honge) ---
print("Loading AI Models... Please wait...")
depth_estimator = pipeline("depth-estimation", model="Intel/dpt-large") # High quality depth model
engine = SculptorEngine()
print("Models Loaded Successfully!")

# --- 2. Core Processing Functions ---

def process_full_stack(input_image, enable_3d):
    if input_image is None:
        return None, None, None, "Please upload an image first."
    
    try:
        # Temporary save
        input_path = "/tmp/input_img.png"
        input_image.save(input_path)
        
        status_msg = "🔧 Step 1/3: Cleaning & Tracing Vectors..."
        yield None, None, None, status_msg
        
        # Vector Generation
        svg_path = engine.generate_vector(input_path)
        
        status_msg = "⚙️ Step 2/3: Converting to DXF (CNC Ready)..."
        yield svg_path, None, None, status_msg
        
        # DXF Generation
        dxf_path = engine.generate_dxf(svg_path)
        
        stl_path = None
        if enable_3d:
            status_msg = "🧠 Step 3/3: Generating 3D Relief (AI Processing)..."
            yield svg_path, dxf_path, None, status_msg
            
            # 3D Generation
            pil_img = Image.open(input_path).convert("RGB")
            stl_path = engine.generate_3d_relief(pil_img, depth_estimator)
        
        status_msg = "✅ Done! Files ready for download."
        yield svg_path, dxf_path, stl_path, status_msg

    except Exception as e:
        yield None, None, None, f"Error: {str(e)}"

# --- 3. UI Design (Dark Theme, Professional) ---

custom_css = """
.gradio-container {
    font-family: 'Inter', sans-serif;
    background: #111827 !important;
    color: white;
}
footer {
    display: none !important;
}
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="cyan")) as demo:
    
    gr.Markdown(
        """
        <div align="center">
        <h1 style="font-size: 3em; font-weight: 800; color: #06b6d4;">Sculptor AI Pro</h1>
        <p style="color: #9ca3af;">Convert Images to Professional CNC Vectors & 3D Reliefs</p>
        </div>
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(label="Upload Source Image", type="pil", height=400)
            
            with gr.Accordion("⚙️ Config", open=False):
                enable_3d = gr.Checkbox(label="Enable 3D Relief (STL) [Requires GPU]", value=False)
                
            submit_btn = gr.Button("🚀 Generate Artwork", variant="primary", size="lg", scale=2)
            
        with gr.Column(scale=1):
            status = gr.Textbox(label="Status", interactive=False, value="Waiting...")
            
            gr.Markdown("### 📂 Download Results")
            out_svg = gr.File(label="Vector (SVG) - For Editing")
            out_dxf = gr.File(label="CNC (DXF) - For Laser/Cutter")
            out_stl = gr.File(label="3D Relief (STL) - For CNC Router")

    # Event Handler with Progress Simulation
    submit_btn.click(
        fn=process_full_stack,
        inputs=[input_image, enable_3d],
        outputs=[out_svg, out_dxf, out_stl, status],
        show_progress="minimal"
    )

if __name__ == "__main__":
    demo.launch()