import gradio as gr
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
import torch
from PIL import Image
import cadquery as cq
from cadquery import exporters
import tempfile
import os
import numpy as np

# Set Hugging Face token (set in Space Secrets as HF_TOKEN)
HF_TOKEN = os.getenv("HF_TOKEN", None)

# Load 2D models
model_id = "Osama03/Finetuned_diffusion_interiordesign"
try:
    txt2img_pipe = StableDiffusionPipeline.from_pretrained(
        model_id, torch_dtype=torch.float16, use_auth_token=HF_TOKEN
    )
    txt2img_pipe = txt2img_pipe.to("cuda" if torch.cuda.is_available() else "cpu")
    img2img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
        model_id, torch_dtype=torch.float16, use_auth_token=HF_TOKEN
    )
    img2img_pipe = img2img_pipe.to("cuda" if torch.cuda.is_available() else "cpu")
except Exception as e:
    print(f"Failed to load {model_id}: {e}. Check token or model availability.")
    raise

# Try to load TripoSR
try:
    from tsr.system import TSR
    triposr_model = TSR.from_pretrained("stabilityai/TripoSR", config_name="config.yaml", weight_name="model.ckpt")
except Exception as e:
    print(f"TripoSR loading failed: {e}. Falling back to CadQuery for 3D.")
    triposr_model = None

def generate_2d_image(prompt, style, features, input_image):
    features_str = ", ".join(features) if features else "various modern elements"
    full_prompt = f"A unique interior design: {prompt}, in {style} style, featuring {features_str}, high detail, realistic lighting, 4k resolution"
    
    if input_image is not None:
        input_image = Image.fromarray(input_image).resize((512, 512))
        image = img2img_pipe(full_prompt, image=input_image, strength=0.75, num_inference_steps=50, guidance_scale=7.5).images[0]
    else:
        image = txt2img_pipe(full_prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
    
    return image

def generate_3d_model(input_image, prompt):
    if triposr_model is not None and input_image is not None:
        try:
            temp_image = Image.fromarray(input_image)
            with tempfile.TemporaryDirectory() as tmpdir:
                triposr_model(temp_image, device="cuda" if torch.cuda.is_available() else "cpu", export_format="glb", output_dir=tmpdir)
                glb_path = os.path.join(tmpdir, "mesh.glb")
                return glb_path
        except Exception as e:
            print(f"TripoSR failed: {e}. Using CadQuery fallback.")
    
    # Fallback: Simple CadQuery room
    room = cq.Workplane("XY").box(5, 4, 3).faces(">Z").shell(-0.1)
    with tempfile.NamedTemporaryFile(suffix=".glb", delete=False) as tmpfile:
        exporters.export(room, tmpfile.name, exportType='GLB')
        return tmpfile.name

def generate_cad_model(room_length, room_width, room_height, features):
    room = cq.Workplane("XY").box(room_length, room_width, room_height).faces(">Z").shell(-0.1)
    feature_positions = {"sofa": (room_length/4, room_width/2, 0.5), "coffee table": (room_length/2, room_width/2, 0.3)}
    feature_sizes = {"sofa": (2, 1, 0.8), "coffee table": (1, 1, 0.4)}
    
    for feat in features:
        if feat in feature_positions:
            pos = feature_positions[feat]
            size = feature_sizes[feat]
            room = room.union(cq.Workplane("XY").transformed(offset=(pos[0], pos[1], pos[2])).box(*size))
    
    with tempfile.NamedTemporaryFile(suffix=".glb", delete=False) as tmpfile:
        exporters.export(room, tmpfile.name, exportType='GLB')
        return tmpfile.name

# UI with Tabs
with gr.Blocks(title="Interior Design Generator MVP") as demo:
    gr.Markdown("# Interior Design Image Generator MVP with 3D & CAD")
    
    with gr.Tabs():
        with gr.Tab("2D Generation"):
            with gr.Row():
                prompt = gr.Textbox(label="Text Prompt")
                style = gr.Dropdown(choices=["modern", "vintage", "minimalist", "industrial", "bohemian", "scandinavian", "rustic"], value="modern")
                features = gr.Checkboxgroup(choices=["sofa", "coffee table", "lamp", "bookshelf", "fireplace", "plants", "artwork", "rug"])
            input_image = gr.Image(label="Upload Input Image (optional)", type="numpy")
            generate_2d_btn = gr.Button("Generate 2D Image")
            output_2d = gr.Image(label="Generated 2D Design")
            generate_2d_btn.click(generate_2d_image, inputs=[prompt, style, features, input_image], outputs=output_2d)
        
        with gr.Tab("3D Visualization"):
            gr.Markdown("Generate and view 3D from the 2D image (or prompt if no image).")
            input_3d_image = gr.Image(label="Use Generated 2D Image (paste from above or upload)", type="numpy")
            input_3d_prompt = gr.Textbox(label="Fallback Prompt (if no image)")
            generate_3d_btn = gr.Button("Generate 3D Model")
            output_3d = gr.Model3D(label="3D Model (rotate/zoom in browser)")
            generate_3d_btn.click(generate_3d_model, inputs=[input_3d_image, input_3d_prompt], outputs=output_3d)
        
        with gr.Tab("CAD-like Editing"):
            gr.Markdown("Parametric editing: Adjust room dimensions and features, regenerate CAD model.")
            room_length = gr.Slider(3, 10, value=5, label="Room Length (m)")
            room_width = gr.Slider(3, 10, value=4, label="Room Width (m)")
            room_height = gr.Slider(2, 4, value=3, label="Room Height (m)")
            cad_features = gr.Checkboxgroup(choices=["sofa", "coffee table"], label="Features (positioned simply)")
            generate_cad_btn = gr.Button("Generate CAD Model")
            output_cad = gr.Model3D(label="CAD 3D Model (edit params to update)")
            generate_cad_btn.click(generate_cad_model, inputs=[room_length, room_width, room_height, cad_features], outputs=output_cad)

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

def generate_3d_model(input_image, prompt):
    if triposr_model is not None and input_image is not None:
        # Use TripoSR for 3D
        try:
            temp_image = Image.fromarray(input_image)
            with tempfile.TemporaryDirectory() as tmpdir:
                triposr_model(temp_image, device="cuda" if torch.cuda.is_available() else "cpu", export_format="glb", output_dir=tmpdir)
                glb_path = os.path.join(tmpdir, "mesh.glb")
                return glb_path
        except Exception as e:
            print(f"TripoSR failed: {e}. Using CadQuery fallback.")
    
    # Fallback: Generate simple 3D room with CadQuery
    room = cq.Workplane("XY").box(5, 4, 3).faces(">Z").shell(-0.1)  # Basic room
    with tempfile.NamedTemporaryFile(suffix=".glb", delete=False) as tmpfile:
        exporters.export(room, tmpfile.name, exportType='GLB')
        return tmpfile.name

def generate_cad_model(room_length, room_width, room_height, features):
    room = cq.Workplane("XY").box(room_length, room_width, room_height).faces(">Z").shell(-0.1)
    feature_positions = {"sofa": (room_length/4, room_width/2, 0.5), "coffee table": (room_length/2, room_width/2, 0.3)}
    feature_sizes = {"sofa": (2, 1, 0.8), "coffee table": (1, 1, 0.4)}
    
    for feat in features:
        if feat in feature_positions:
            pos = feature_positions[feat]
            size = feature_sizes[feat]
            room = room.union(cq.Workplane("XY").transformed(offset=(pos[0], pos[1], pos[2])).box(*size))
    
    with tempfile.NamedTemporaryFile(suffix=".glb", delete=False) as tmpfile:
        exporters.export(room, tmpfile.name, exportType='GLB')
        return tmpfile.name

# UI with Tabs
with gr.Blocks(title="Interior Design Generator MVP") as demo:
    gr.Markdown("# Interior Design Image Generator MVP with 3D & CAD")
    
    with gr.Tabs():
        with gr.Tab("2D Generation"):
            with gr.Row():
                prompt = gr.Textbox(label="Text Prompt")
                style = gr.Dropdown(choices=["modern", "vintage", "minimalist", "industrial", "bohemian", "scandinavian", "rustic"], value="modern")
                features = gr.Checkboxgroup(choices=["sofa", "coffee table", "lamp", "bookshelf", "fireplace", "plants", "artwork", "rug"])
            input_image = gr.Image(label="Upload Input Image (optional)", type="numpy")
            generate_2d_btn = gr.Button("Generate 2D Image")
            output_2d = gr.Image(label="Generated 2D Design")
            generate_2d_btn.click(generate_2d_image, inputs=[prompt, style, features, input_image], outputs=output_2d)
        
        with gr.Tab("3D Visualization"):
            gr.Markdown("Generate and view 3D from the 2D image (or prompt if no image).")
            input_3d_image = gr.Image(label="Use Generated 2D Image (paste from above or upload)", type="numpy")
            input_3d_prompt = gr.Textbox(label="Fallback Prompt (if no image)")
            generate_3d_btn = gr.Button("Generate 3D Model")
            output_3d = gr.Model3D(label="3D Model (rotate/zoom in browser)")
            generate_3d_btn.click(generate_3d_model, inputs=[input_3d_image, input_3d_prompt], outputs=output_3d)
        
        with gr.Tab("CAD-like Editing"):
            gr.Markdown("Parametric editing: Adjust room dimensions and features, regenerate CAD model.")
            room_length = gr.Slider(3, 10, value=5, label="Room Length (m)")
            room_width = gr.Slider(3, 10, value=4, label="Room Width (m)")
            room_height = gr.Slider(2, 4, value=3, label="Room Height (m)")
            cad_features = gr.Checkboxgroup(choices=["sofa", "coffee table"], label="Features (positioned simply)")
            generate_cad_btn = gr.Button("Generate CAD Model")
            output_cad = gr.Model3D(label="CAD 3D Model (edit params to update)")
            generate_cad_btn.click(generate_cad_model, inputs=[room_length, room_width, room_height, cad_features], outputs=output_cad)

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