app.py

from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks
from fastapi.responses import FileResponse, JSONResponse
from uuid import uuid4
from pathlib import Path
import shutil
import os
import json
from dotenv import load_dotenv
from tasks import process_image_task

load_dotenv()

# Directories
# Use /tmp for Hugging Face Spaces as it is writable
UPLOAD_DIR = Path("/tmp/uploads")
RESULT_DIR = Path("/tmp/results")
UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
RESULT_DIR.mkdir(parents=True, exist_ok=True)

# In-memory job store (Global variable)
# Since HF Spaces (Free) runs 1 replica, this works for a demo.
JOBS = {}

app = FastAPI(title="Depth->STL processing service (Standalone)")

def update_job_status(job_id: str, state: str, detail: str = "", result: str = ""):
    JOBS[job_id] = {
        "state": state,
        "detail": detail,
        "result": result
    }

@app.get("/health")
def health_check():
    return {"status": "ok"}

@app.post("/upload/")
async def upload_image(background_tasks: BackgroundTasks, file: UploadFile = File(...)):
    # Basic validation
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="File must be an image")

    job_id = str(uuid4())
    safe_name = Path(file.filename).name
    fname = f"{job_id}_{safe_name}"
    save_path = UPLOAD_DIR / fname

    # Save uploaded file
    try:
        with save_path.open("wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to save upload: {e}")

    # Mark queued
    update_job_status(job_id, "QUEUED", "Job received and queued")

    # Add to background tasks
    background_tasks.add_task(
        process_image_task,
        str(save_path),
        str(RESULT_DIR),
        job_id,
        update_job_status
    )

    return {"job_id": job_id}


@app.get("/status/{job_id}")
def status(job_id: str):
    job = JOBS.get(job_id)
    if not job:
        return JSONResponse({"state": "UNKNOWN", "detail": "No such job_id"}, status_code=404)
    return JSONResponse(job)


@app.get("/download/{job_id}")
def download(job_id: str):
    print(f"[DEBUG] Download request for {job_id}")
    job = JOBS.get(job_id)
    if not job:
        print(f"[DEBUG] Job {job_id} not found in JOBS: {list(JOBS.keys())}")
        raise HTTPException(status_code=404, detail="No such job")
    
    if job.get("state") != "SUCCESS":
        print(f"[DEBUG] Job {job_id} state is {job.get('state')}")
        raise HTTPException(status_code=404, detail="Result not ready")
    
    stl_path = job.get("result")
    print(f"[DEBUG] Checking file path: {stl_path}")
    if not stl_path or not Path(stl_path).exists():
        print(f"[DEBUG] File missing at {stl_path}")
        # List dir to see what's there
        try:
            parent = Path(stl_path).parent
            print(f"[DEBUG] Contents of {parent}: {list(parent.glob('*'))}")
        except Exception as e:
            print(f"[DEBUG] Failed to list dir: {e}")
            
        raise HTTPException(status_code=404, detail=f"Result file missing at {stl_path}")
    
    return FileResponse(path=stl_path, filename=Path(stl_path).name, media_type="application/sla")