import gradio as gr
import spaces
import torch
from diffusers import DiffusionPipeline
import numpy as np
from PIL import Image
import os
from rembg import remove

# --- MODELS LOADING (Hugging Face Cache se load honge) ---
print("Loading FLUX (Text-to-Image) model...")
# FLUX.1-schnell use kar rahe hain speed ke liye
txt2img_pipe = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", 
    torch_dtype=torch_dtype,
    device_map="auto"
)

print("Loading Image-to-3D model...")
# Hum standard shape generation pipeline loading import use karenge
# (Hugging Face par different Space backends TripoSR ko different tarah se import karte hain, 
# free Tier ke liye yeh standard approach hai)
try:
    from tsr.system import TSR
    from tsr.utils import remove_background, resize_foreground, save_video
    IS_TSR_AVAILABLE = True
    # TSR Model Load (CPU par initialize, GPU decorator switch karega)
    tsr_model = TSR.from_pretrained(
        "stabilityai/TripoSR",
        config_name="config.yaml",
        weight_name="model.ckpt"
    )
except ImportError:
    print("Warning: TripoSR libraries not found. Image-to-3D will be dummy.")
    IS_TSR_AVAILABLE = False


# --- HELPERS ---

def process_image_background(image_path):
    """Image se background hatata hai 3D generation se pehle"""
    input_image = Image.open(image_path)
    # Using rembg for background removal
    output_image = remove(input_image)
    # Save as PNG with alpha channel
    bg_removed_path = "processed_input.png"
    output_image.save(bg_removed_path)
    return bg_removed_path

def generate_3d_from_processed_image(processed_image_path):
    """TripoSR ka use karke PNG se .glb banata hai"""
    if not IS_TSR_AVAILABLE:
        return None # Backend error safety

    # Run TripoSR generation
    # ZeroGPU handle karega move to device automatically decorator ke sath
    scene_codes = tsr_model(processed_image_path, device="cuda")
    
    # Mesh extract karke .glb format mein save karna
    # (Extracting vertices/faces internally and saving)
    meshes = tsr_model.extract_mesh(scene_codes)
    glb_path = "generated_model.glb"
    meshes[0].export(glb_path) # Export first generated mesh
    
    return glb_path


# --- GENERATION FUNCTIONS (with GPU Decorator) ---

@spaces.GPU(duration=60) # Text to Image thoda time leta hai
def text_to_3d_pipeline(prompt):
    if not prompt:
        return None, None
        
    print(f"Generating Image for prompt: {prompt}")
    # 1. Text to Real Image (FLUX)
    # inference_steps kam rakhe hain speed ke liye
    image_result = txt2img_pipe(
        prompt, 
        guidance_scale=0.0, 
        num_inference_steps=4, 
        max_sequence_length=256
    ).images[0]
    
    temp_img_path = "text2img_output.png"
    image_result.save(temp_img_path)
    
    print("Generating 3D model from image...")
    # 2. Background Removal & 3D Generation
    if IS_TSR_AVAILABLE:
        processed_img = process_image_background(temp_img_path)
        glb_file = generate_3d_from_processed_image(processed_img)
        return temp_img_path, glb_file
    else:
        return temp_img_path, None # Image dikhegi par 3D nahi banega safety ke liye


@spaces.GPU(duration=30) # Image to 3D faster hota hai
def image_to_3d_pipeline(image_filepath):
    if not image_filepath or not IS_TSR_AVAILABLE:
        return None
        
    print("Processing uploaded image for 3D...")
    # 1. Background Removal
    processed_img = process_image_background(image_filepath)
    
    # 2. 3D Generation
    glb_file = generate_3d_from_processed_image(processed_img)
    
    return glb_file


# --- UI DESIGN ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🌌 V.O.I.D - Real 3D Generation Studio")
    gr.Markdown("Generating real images and 3D models using FLUX.1 & TripoSR (ZeroGPU).")
    
    with gr.Row():
        # Left Side: Input Options
        with gr.Column(scale=1):
            with gr.Tabs():
                # Tab 1: Text to 3D
                with gr.TabItem("📝 Text to 3D Pipeline"):
                    txt_prompt = gr.Textbox(
                        label="Describe your 3D model", 
                        placeholder="e.g., A photorealistic golden crown, a cyberpunk helmet..."
                    )
                    txt_btn = gr.Button("Generate Image & 3D 🚀", variant="primary")
                    # Intermediate Image Output
                    gen_img_output = gr.Image(label="Generated Reference Image", type="filepath")
                
                # Tab 2: Image to 3D
                with gr.TabItem("🖼️ Image to 3D"):
                    img_input = gr.Image(type="filepath", label="Upload Photo (Background will be removed)")
                    img_btn = gr.Button("Generate 3D Model 🚀", variant="primary")
            
        # Right Side: Final 3D Model Viewer
        with gr.Column(scale=1):
            model_viewer = gr.Model3D(
                label="Final 3D Model Viewer (.glb)", 
                height=500
            )
            
    # --- LOGIC CONNECTIONS ---
    
    # Text-to-3D returns [Image, GLB]
    txt_btn.click(
        fn=text_to_3d_pipeline, 
        inputs=[txt_prompt], 
        outputs=[gen_img_output, model_viewer]
    )
    
    # Image-to-3D returns [GLB]
    img_btn.click(
        fn=image_to_3d_pipeline, 
        inputs=[img_input], 
        outputs=[model_viewer]
    )

# Launch
demo.launch()