import os
import sys
import uuid
import io
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import FileResponse
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image

# Route Python to the cloned TripoSR repository
sys.path.append("/app/TripoSR")
import torch
from tsr.system import TSR
from tsr.utils import remove_background, resize_foreground

app = FastAPI(title="TripoSR CPU Free API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

model = None
OUTPUT_DIR = "/app/outputs"

@app.on_event("startup")
async def startup_event():
    global model
    print("Loading TripoSR model to CPU...")
    try:
        # Load the model and explicitly map it to the CPU
        model = TSR.from_pretrained(
            "stabilityai/TripoSR",
            config_name="config.yaml",
            weight_name="model.ckpt"
        )
        # Lower chunk size to protect the 16GB RAM limit
        model.renderer.set_chunk_size(131072) 
        model.to("cpu")
        print("TripoSR successfully loaded and running on CPU!")
    except Exception as e:
        print(f"Error loading model: {e}")

@app.post("/generate-3d")
async def generate_3d_model(image: UploadFile = File(...)):
    """
    Endpoint that accepts 1 image, removes the background, and generates a 3D model.
    """
    if model is None:
        raise HTTPException(status_code=503, detail="Model is still loading.")

    try:
        # 1. Read the uploaded image
        contents = await image.read()
        pil_img = Image.open(io.BytesIO(contents)).convert("RGB")
        
        # 2. Pre-process: Remove background and center the object
        print("Removing background...")
        img_nobg = remove_background(pil_img, rembg_session=None)
        img_resized = resize_foreground(img_nobg, 0.85)
        
        # 3. Generate 3D Model (This takes 2-5 minutes on a CPU)
        print("Generating 3D model on CPU...")
        with torch.no_grad():
            scene_codes = model.get_scene_codes(img_resized)
            mesh = model.extract_mesh(scene_codes)[0]
        
        # 4. Save and return the file
        job_id = str(uuid.uuid4())
        output_path = os.path.join(OUTPUT_DIR, f"{job_id}.obj")
        
        mesh.export(output_path)
        
        return FileResponse(
            path=output_path, 
            media_type="application/octet-stream", 
            filename=f"{job_id}.obj"
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")