backend/app/main.py

"""FastAPI application: routes, CORS, static mounts.

Endpoints (spec section 9):
  POST /upload-room    - accept 3-10 images, preprocess, store, return ids
  POST /analyze-room   - detect + floor/usable polygon + free space for one image
  GET  /get-catalog    - catalog JSON, filterable, recommendation-ranked
  POST /generate-room  - synchronous placement + inpaint, returns output image URL

There is intentionally NO /remove-object endpoint (object removal is OUT).
"""
from __future__ import annotations

import uuid
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import Depends, FastAPI, File, Form, HTTPException, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from sqlalchemy.orm import Session

from . import pipeline
from .config import settings
from .db import get_db, init_db
from .models import Generation, RoomImage, UploadSession
from .modules import recommend
from .modules.input_preprocess import preprocess_image
from .modules.ml_runtime import ml_available
from .schemas import (
    AnalyzeResponse,
    CatalogItem,
    DetectedObject,
    GenerateResponse,
    ImageMeta,
    Placement,
    UploadResponse,
)

ALLOWED_TYPES = {"image/jpeg", "image/jpg", "image/png", "image/webp"}


@asynccontextmanager
async def lifespan(app: FastAPI):
    init_db()
    yield


app = FastAPI(title="Room Visualizer AI", version="1.0.0", lifespan=lifespan)

if settings.CORS_ALLOW_ALL:
    # Dev-friendly: lets Expo web and a phone on the LAN call the API.
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_credentials=False,
        allow_methods=["*"],
        allow_headers=["*"],
    )
else:
    app.add_middleware(
        CORSMiddleware,
        allow_origins=settings.cors_origins,
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )
app.mount("/static", StaticFiles(directory=str(settings.STATIC_DIR)), name="static")


def abs_url(request: Request, rel: str | None) -> str | None:
    if not rel:
        return None
    if rel.startswith("http"):
        return rel
    # Build from the incoming request host so the same backend serves the web app
    # (127.0.0.1) and a phone on the LAN (the dev machine's IP) automatically.
    base = settings.PUBLIC_BASE_URL or str(request.base_url).rstrip("/")
    return f"{base}/{rel.lstrip('/')}"


# ----- request bodies -------------------------------------------------------
class AnalyzeRequest(BaseModel):
    session_id: str
    image_id: str


class GenerateRequest(BaseModel):
    session_id: str
    image_id: str
    item_ids: list[str]
    hints: dict | None = None


# ----- routes ---------------------------------------------------------------
@app.get("/")
def root():
    return {
        "status": "ok",
        "service": "room-visualizer-ai",
        "provider": settings.IMAGE_PROVIDER,
        "ml_available": ml_available(),
    }


@app.post("/upload-room", response_model=UploadResponse)
def upload_room(
    request: Request,
    files: list[UploadFile] = File(...),
    session_id: str | None = Form(None),
    db: Session = Depends(get_db),
) -> UploadResponse:
    # Supports two modes against the same endpoint:
    #  - one-shot: the web app sends all files at once (no session_id);
    #  - incremental: the mobile app uses the native uploader (one file per
    #    request) and passes session_id to append to the same room.
    if session_id:
        session = db.get(UploadSession, session_id)
        if session is None:
            raise HTTPException(404, "Upload session not found.")
    else:
        session = UploadSession()
        db.add(session)
        db.flush()  # populate session.id

    if len(session.images) + len(files) > settings.MAX_IMAGES:
        raise HTTPException(400, f"Too many images (max {settings.MAX_IMAGES} per room).")

    for f in files:
        data = f.file.read()
        if len(data) > settings.MAX_FILE_MB * 1024 * 1024:
            raise HTTPException(400, f"{f.filename} exceeds {settings.MAX_FILE_MB} MB limit.")

        img_id = uuid.uuid4().hex
        raw_path = settings.UPLOAD_DIR / f"{img_id}_raw"
        out_path = settings.UPLOAD_DIR / f"{img_id}.jpg"
        raw_path.write_bytes(data)
        try:
            # preprocess_image validates the bytes are a decodable image.
            res = preprocess_image(raw_path, out_path)
        except ValueError as exc:
            raw_path.unlink(missing_ok=True)
            raise HTTPException(400, str(exc)) from exc
        finally:
            raw_path.unlink(missing_ok=True)

        db.add(
            RoomImage(
                id=img_id,
                session_id=session.id,
                filename=f.filename or f"{img_id}.jpg",
                path=str(out_path),
                rel_path=res.rel_path,
                width=res.width,
                height=res.height,
            )
        )

    db.commit()
    images = [
        ImageMeta(id=im.id, url=abs_url(request, im.rel_path), width=im.width, height=im.height)
        for im in session.images
    ]
    return UploadResponse(session_id=session.id, images=images)


@app.post("/analyze-room", response_model=AnalyzeResponse)
def analyze_room(
    req: AnalyzeRequest, request: Request, db: Session = Depends(get_db)
) -> AnalyzeResponse:
    rec = db.get(RoomImage, req.image_id)
    if rec is None or rec.session_id != req.session_id:
        raise HTTPException(404, "Image not found for this session.")

    result = pipeline.analyze_room(Path(rec.path))
    return AnalyzeResponse(
        session_id=req.session_id,
        image_id=req.image_id,
        width=result.width,
        height=result.height,
        detected_objects=[
            DetectedObject(label=d.label, confidence=round(d.confidence, 3), box=d.box)
            for d in result.detected_objects
        ],
        floor_polygon=result.usable_polygon,
        blocked_areas=result.blocked_areas,
        free_space_ratio=round(result.free_ratio, 4),
        depth_url=abs_url(request, result.depth_rel_path),
        ml_used=ml_available(),
    )


@app.get("/get-catalog", response_model=list[CatalogItem])
def get_catalog(
    request: Request,
    category: str | None = None,
    styles: str | None = None,
) -> list[CatalogItem]:
    style_list = [s.strip() for s in styles.split(",") if s.strip()] if styles else []
    items = recommend.filter_and_rank(category=category, styles=style_list)
    return [
        CatalogItem(
            id=it["id"],
            name=it["name"],
            category=it["category"],
            dimensions=it["dimensions"],
            price=it["price"],
            image_url=abs_url(request, it["image_url"]),
            style_tags=it["style_tags"],
            score=it.get("score"),
        )
        for it in items
    ]


@app.post("/generate-room", response_model=GenerateResponse)
def generate_room(
    req: GenerateRequest, request: Request, db: Session = Depends(get_db)
) -> GenerateResponse:
    rec = db.get(RoomImage, req.image_id)
    if rec is None or rec.session_id != req.session_id:
        raise HTTPException(404, "Image not found for this session.")
    if not req.item_ids:
        raise HTTPException(400, "Select at least one catalog item.")

    catalog = {c["id"]: c for c in recommend.load_catalog()}
    items = [catalog[i] for i in req.item_ids if i in catalog]
    if not items:
        raise HTTPException(400, "No valid catalog items selected.")

    result = pipeline.generate_room(Path(rec.path), items, req.hints)

    gen = Generation(
        session_id=req.session_id,
        image_id=req.image_id,
        item_ids=",".join(req.item_ids),
        provider=result["provider"],
        output_path=result["output_path"],
        output_rel_path=result["output_rel_path"],
    )
    db.add(gen)
    db.commit()

    return GenerateResponse(
        generation_id=gen.id,
        image_url=abs_url(request, result["output_rel_path"]),
        provider=result["provider"],
        placements=[
            Placement(item_id=p.item_id, footprint_px=p.footprint_px, scale_note=p.scale_note)
            for p in result["placements"]
        ],
    )