from fastapi import FastAPI, UploadFile, File, HTTPException, Query
from fastapi.responses import FileResponse
import os
import shutil
import uuid
import requests
from typing import Optional
from app.utils import run_inference

from huggingface_hub import login

hf_token = os.environ.get("HF_TOKEN")
print(hf_token)
login(token=hf_token)


app = FastAPI(title="Stable Fast 3D API")

@app.get("/")
async def root():
    return {"message": "Welcome to Stable Fast 3D API. Use /generate-3d endpoint to convert 2D images to 3D models."}

@app.post("/generate-3d/")
async def generate_3d_model_upload(image: UploadFile = File(...)):
    """Generate 3D model from uploaded image file"""
    return await process_image(image=image)

@app.get("/generate-3d/")
async def generate_3d_model_url(image_url: str = Query(..., description="URL of the image to convert to 3D")):
    """Generate 3D model from image URL"""
    return await process_image(image_url=image_url)

async def process_image(image: Optional[UploadFile] = None, image_url: Optional[str] = None):
    # Create unique ID for this request
    temp_id = str(uuid.uuid4())
    input_path = f"/app/tmp/{temp_id}.png"
    output_dir = f"/app/tmp/{temp_id}_output"
    os.makedirs(output_dir, exist_ok=True)

    try:
        # Handle image from upload or URL
        if image:
            with open(input_path, "wb") as f:
                shutil.copyfileobj(image.file, f)
        elif image_url:
            response = requests.get(image_url, stream=True)
            if response.status_code != 200:
                raise HTTPException(status_code=400, detail="Could not download image from URL")
                
            with open(input_path, "wb") as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)
        else:
            raise HTTPException(status_code=400, detail="Either image file or image URL must be provided")

        # Run the inference
        glb_path = run_inference(input_path, output_dir)
        
        if not os.path.exists(glb_path):
            raise HTTPException(status_code=500, detail="Failed to generate 3D model")
            
        # Return the GLB file
        return FileResponse(
            path=glb_path, 
            media_type="model/gltf-binary", 
            filename="model.glb",
            headers={"Content-Disposition": f"attachment; filename=model.glb"}
        )
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
    
    finally:
        # Clean up temporary files
        if os.path.exists(input_path):
            os.remove(input_path)