from fastapi import FastAPI, Depends, HTTPException, Header, Form
import tempfile
import os
import base64
import json

import depth_pro

app = FastAPI()

# Constants
API_KEY = "mysecureapikey"

# Load models once
depth_model, transform = depth_pro.create_model_and_transforms()
depth_model.eval()

# API Key verification
def verify_key(x_api_key: str = Header(...)):
    if x_api_key != API_KEY:
        raise HTTPException(status_code=403, detail="Invalid API Key")

@app.post("/depth")
async def calculate_depth(
    detections: str = Form(...),  # Accept as JSON string
    image_base64: str = Form(...),
    _: str = Depends(verify_key)
):
    # Decode detection string into Python object
    try:
        detections = json.loads(detections)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Invalid JSON in 'detections': {e}")

    # Decode and temporarily save the image
    try:
        image_bytes = base64.b64decode(image_base64)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Image base64 decode failed: {e}")

    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp:
        tmp.write(image_bytes)
        tmp_path = tmp.name

    try:
        # Load image and run depth estimation
        image_depth, _, f_px = depth_pro.load_rgb(tmp_path)
        depth_input = transform(image_depth)
        prediction = depth_model.infer(depth_input, f_px=f_px)
        depth_np = prediction["depth"].squeeze().cpu().numpy()
    finally:
        os.remove(tmp_path)  # Cleanup temp file

    # Process each detection
    result_data = []
    for item in detections:
        try:
            x1, y1, x2, y2 = item['bounding_box'].values()
            center_x = (x1 + x2) // 2
            center_y = (y1 + y2) // 2
            distance = float(depth_np[center_y][center_x])
            result_data.append({
                "object": item['class_name'],
                "bounding_box": item['bounding_box'],
                "center": {"x": center_x, "y": center_y},
                "distance_m": round(distance, 2)
            })
        except Exception as e:
            result_data.append({
                "object": item.get('class_name', 'unknown'),
                "error": f"Failed to compute distance: {e}"
            })

    return {"depth_results": result_data}