main.py

"""
main.py — FastAPI REST API for Reunite
───────────────────────────────────────────────────────────────────────────────
Start:  uvicorn main:app --reload --host 0.0.0.0 --port 8000
Docs:   http://localhost:8000/docs

Route map:
  GET   /api/health                   — liveness probe + registry count
  GET   /api/stats                    — totals: registered / missing / found
  POST  /api/register                 — register a new missing person
  POST  /api/search                   — face-similarity search
  GET   /api/persons                  — list all registered persons
  GET   /api/persons/{id}/image       — serve the stored photo
  PATCH /api/persons/{id}/status      — mark as missing / found
  DELETE /api/persons/{id}            — remove from registry

  /  → served from ../frontend/ (StaticFiles, html=True)
───────────────────────────────────────────────────────────────────────────────
"""

from __future__ import annotations

from pathlib import Path
import logging

from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles

import core

logger = logging.getLogger("uvicorn.error")

# ─────────────────────────────────────────────────────────────────────────────
app = FastAPI(
    title="Reunite — Missing Persons AI",
    version="2.0.0",
    docs_url="/docs",
    redoc_url="/redoc",
)

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

PERSONS_DIR = Path(__file__).parent / "Storage" / "persons"

_frontend_root = Path(__file__).parent.parent
FRONTEND_DIR = (
    _frontend_root / "Frontend"
    if (_frontend_root / "Frontend").exists()
    else (_frontend_root / "frontend")
)


# ─────────────────────────────────────────────────────────────────────────────
@app.on_event("startup")
async def on_startup():
    core.ensure_dirs()
    # Warm up model on CPU so first request is faster (weights may still download).
    try:
        core.get_face_app()
        logger.info("ArcFace model initialized (CPU).")
    except Exception as e:
        logger.exception("ArcFace model warmup failed: %s", e)


# ─────────────────────────────────────────────────────────────────────────────
# Health & Stats
# ─────────────────────────────────────────────────────────────────────────────

@app.get("/api/health", tags=["meta"])
async def health():
    persons = core.get_all_persons()
    return {"status": "ok", "registered": len(persons)}


@app.get("/api/stats", tags=["meta"])
async def stats():
    persons = core.get_all_persons()
    total   = len(persons)
    found   = sum(1 for p in persons if p.get("status") == "found")
    return {
        "total":   total,
        "missing": total - found,
        "found":   found,
    }


# ─────────────────────────────────────────────────────────────────────────────
# Register a missing person
# ─────────────────────────────────────────────────────────────────────────────

@app.post("/api/register", tags=["persons"])
async def register(
    image:               UploadFile = File(...,  description="Face photo (JPG / PNG / WEBP)"),
    name:                str        = Form(...,  description="Full name"),
    phone_contact:       str        = Form(...,  description="Contact phone number"),
    age:                 str        = Form("",   description="Approximate age"),
    gender:              str        = Form("",   description="male | female | other"),
    last_seen_date:      str        = Form("",   description="YYYY-MM-DD"),
    last_seen_location:  str        = Form("",   description="City / area last seen"),
    address:             str        = Form("",   description="Home address"),
    national_id:         str        = Form("",   description="National ID (optional)"),
    description:         str        = Form("",   description="Additional notes"),
):
    image_bytes = await image.read()
    result = core.register_missing_person(
        image_bytes,
        {
            "name":               name,
            "age":                age,
            "gender":             gender,
            "last_seen_date":     last_seen_date,
            "last_seen_location": last_seen_location,
            "phone_contact":      phone_contact,
            "address":            address,
            "national_id":        national_id,
            "description":        description,
        },
    )
    if not result["success"]:
        raise HTTPException(status_code=422, detail=result["error"])
    timing = result.get("timing_ms")
    if isinstance(timing, dict):
        logger.info(
            "Registered person %s | embedding=%sms io=%sms total=%sms",
            result.get("id"),
            timing.get("embedding_ms"),
            timing.get("io_ms"),
            timing.get("total_ms"),
        )
    return result


# ─────────────────────────────────────────────────────────────────────────────
# Search by face image
# ─────────────────────────────────────────────────────────────────────────────

@app.post("/api/search", tags=["search"])
async def search(
    image:     UploadFile = File(...,   description="Query face photo"),
    top_k:     int        = Form(5,     description="Max candidates to return"),
    threshold: float      = Form(0.38,  description="Minimum cosine similarity"),
):
    image_bytes = await image.read()
    result = core.search_person(image_bytes, top_k=top_k, threshold=threshold)
    if not result["success"]:
        raise HTTPException(status_code=422, detail=result["error"])
    return result


# ─────────────────────────────────────────────────────────────────────────────
# Directory
# ─────────────────────────────────────────────────────────────────────────────

@app.get("/api/persons", tags=["persons"])
async def list_persons():
    return {"persons": core.get_all_persons()}


@app.get("/api/persons/{pid}/image", tags=["persons"])
async def person_image(pid: str):
    path = PERSONS_DIR / pid / "photo.jpg"
    if not path.exists():
        raise HTTPException(status_code=404, detail="Image not found.")
    return FileResponse(str(path), media_type="image/jpeg")


@app.patch("/api/persons/{pid}/status", tags=["persons"])
async def update_status(pid: str, status: str = Form(...)):
    if status not in ("missing", "found"):
        raise HTTPException(status_code=400, detail="status must be 'missing' or 'found'.")
    result = core.update_person_status(pid, status)
    if not result["success"]:
        raise HTTPException(status_code=404, detail=result["error"])
    return result


@app.delete("/api/persons/{pid}", tags=["persons"])
async def delete_person(pid: str):
    result = core.delete_person(pid)
    if not result["success"]:
        raise HTTPException(status_code=404, detail=result["error"])
    return result


# ─────────────────────────────────────────────────────────────────────────────
# Frontend — must be mounted LAST so API routes take priority
# ─────────────────────────────────────────────────────────────────────────────
if FRONTEND_DIR.exists():
    app.mount(
        "/",
        StaticFiles(directory=str(FRONTEND_DIR), html=True),
        name="frontend",
    )