"""
Face Recognition Backend API on Hugging Face Spaces (Gradio SDK + FastAPI).
Serves MTCNN + ArcFace vector inference with full CORS support for Vercel.
"""

import base64
import os
import secrets
import sys
import threading
import time
import traceback
from typing import Optional, List, Dict, Any

from fastapi import FastAPI, Request, Response, HTTPException, status
from fastapi.responses import JSONResponse, Response as RawResponse
from fastapi.middleware.cors import CORSMiddleware
from starlette.middleware.base import BaseHTTPMiddleware
from pydantic import BaseModel
import gradio as gr

import pipeline as P

# --------------------------------------------------------------------------
# State & Sessions
# --------------------------------------------------------------------------
SESSION_TTL = 4 * 3600
SESSIONS = {}
_SESS_LOCK = threading.Lock()
_status = {"ready": False, "error": None, "seconds": None}


def _warm():
    t0 = time.time()
    try:
        P.warmup()
        _status["ready"] = True
    except Exception as exc:
        _status["error"] = f"{type(exc).__name__}: {exc}"
        traceback.print_exc()
    finally:
        _status["seconds"] = round(time.time() - t0, 1)
        print(f"[warmup] ready={_status['ready']} in {_status['seconds']}s")


threading.Thread(target=_warm, daemon=True).start()


def _get_session(sid: Optional[str]):
    now = time.time()
    with _SESS_LOCK:
        for dead in [s for s, v in SESSIONS.items() if now - v["touched"] > SESSION_TTL]:
            SESSIONS.pop(dead, None)
        if not sid or sid not in SESSIONS:
            sid = secrets.token_urlsafe(18)
            SESSIONS[sid] = {"gallery": P.Gallery(), "dismissed": set(), "touched": now}
        SESSIONS[sid]["touched"] = now
        return sid, SESSIONS[sid]["gallery"], SESSIONS[sid]["dismissed"]


def _session_payload(gallery, dismissed):
    return {
        "people": gallery.summary(),
        "total_embeddings": sum(len(v) for v in gallery.data.values()),
        "duplicates": P.find_duplicates(gallery.data, dismissed),
    }


# --------------------------------------------------------------------------
# FastAPI Application
# --------------------------------------------------------------------------
api_app = FastAPI(title="Face Recognition API", docs_url="/docs")

# CORS middleware for cross-origin Vercel requests
api_app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["Content-Disposition", "X-Session-Id"],
)


class SessionMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        sid_in = request.cookies.get("sid") or request.headers.get("X-Session-Id")
        sid, gallery, dismissed = _get_session(sid_in)
        request.state.sid = sid
        request.state.gallery = gallery
        request.state.dismissed = dismissed

        response = await call_next(request)

        # Attach session id to cookie & response header
        response.set_cookie(
            "sid", sid, httponly=True, samesite="none", secure=True, max_age=SESSION_TTL
        )
        response.headers["X-Session-Id"] = sid
        return response


api_app.add_middleware(SessionMiddleware)


# --------------------------------------------------------------------------
# Pydantic Request Models
# --------------------------------------------------------------------------
class EnrollRequest(BaseModel):
    person_id: str
    images: List[str]


class RecognizeRequest(BaseModel):
    image: str


class MergeRequest(BaseModel):
    source: str
    target: str


class DismissRequest(BaseModel):
    a: str
    b: str


# --------------------------------------------------------------------------
# REST API Endpoints
# --------------------------------------------------------------------------
@api_app.get("/api/status")
async def get_status(request: Request):
    people_count = len(request.state.gallery.data) if hasattr(request.state, "gallery") else 0
    return {**_status, "people": people_count}


@api_app.get("/api/session")
async def get_session_info(request: Request):
    return _session_payload(request.state.gallery, request.state.dismissed)


@api_app.post("/api/session/reset")
async def reset_session(request: Request):
    with _SESS_LOCK:
        SESSIONS[request.state.sid]["gallery"] = P.Gallery()
        SESSIONS[request.state.sid]["dismissed"] = set()
    return {"ok": True}


@api_app.post("/api/enroll")
async def enroll(req: EnrollRequest, request: Request):
    person_id = req.person_id.strip()
    images = req.images

    if not person_id:
        raise HTTPException(status_code=400, detail="A name is required")
    if not images:
        raise HTTPException(status_code=400, detail="No images supplied")
    if not _status["ready"]:
        raise HTTPException(status_code=503, detail="Models are still loading, please wait")

    results = []
    added = 0
    for idx, data_url in enumerate(images):
        try:
            bgr = P.decode_data_url(data_url)
            faces = P.extract_faces(bgr)
            if not faces:
                results.append({"index": idx, "ok": False, "reason": "No face detected"})
                continue
            if len(faces) > 1:
                results.append({
                    "index": idx, "ok": False,
                    "reason": f"{len(faces)} faces in frame - register one face at a time"
                })
                continue

            face = faces[0]
            request.state.gallery.add(person_id, P.embed(face["crop_bgr"]))
            added += 1
            results.append({
                "index": idx,
                "ok": True,
                "confidence": round(face["confidence"], 4),
                "crop": P.encode_jpeg_data_url(face["crop_bgr"], quality=75),
            })
        except Exception as exc:
            results.append({"index": idx, "ok": False, "reason": str(exc)})

    return {
        "person_id": person_id,
        "added": added,
        "total": len(request.state.gallery.data.get(person_id, [])),
        "cohesion": P.cohesion(request.state.gallery.data.get(person_id, [])),
        "results": results,
    }


@api_app.post("/api/recognize")
async def recognize(req: RecognizeRequest, request: Request):
    data_url = req.image
    if not data_url:
        raise HTTPException(status_code=400, detail="Image is required")
    if not _status["ready"]:
        raise HTTPException(status_code=503, detail="Models are still loading")
    if not request.state.gallery.data:
        raise HTTPException(status_code=400, detail="Nobody registered yet")

    t0 = time.time()
    try:
        bgr = P.decode_data_url(data_url)
        h, w = bgr.shape[:2]
        faces = P.extract_faces(bgr)
    except Exception as exc:
        raise HTTPException(status_code=400, detail=str(exc))

    out = []
    for face in faces:
        ranked = request.state.gallery.rank(P.embed(face["crop_bgr"]))
        top = ranked[0] if ranked else None
        runner = ranked[1] if len(ranked) > 1 else None
        out.append({
            "box": face["box"],
            "confidence": round(face["confidence"], 4),
            "crop": P.encode_jpeg_data_url(face["crop_bgr"], quality=70),
            "ranked": [{"id": r["id"], "score": round(r["score"], 4)} for r in ranked],
            "top": {"id": top["id"], "score": round(top["score"], 4)} if top else None,
            "margin": round(top["score"] - runner["score"], 4) if runner else None,
        })

    return {
        "faces": out,
        "frame": {"width": w, "height": h},
        "elapsed_ms": int((time.time() - t0) * 1000),
    }


@api_app.delete("/api/person/{person_id}")
async def delete_person(person_id: str, request: Request):
    ok = request.state.gallery.remove(person_id)
    for pair in [p for p in request.state.dismissed if person_id in p]:
        request.state.dismissed.discard(pair)
    return {"ok": ok}


@api_app.post("/api/person/merge")
async def merge_people(req: MergeRequest, request: Request):
    source, target = req.source.strip(), req.target.strip()
    try:
        request.state.gallery.merge(source, target)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc))
    for pair in [p for p in request.state.dismissed if source in p]:
        request.state.dismissed.discard(pair)
    return {"ok": True, **_session_payload(request.state.gallery, request.state.dismissed)}


@api_app.post("/api/duplicates/dismiss")
async def dismiss_duplicate(req: DismissRequest, request: Request):
    a, b = req.a.strip(), req.b.strip()
    gal = request.state.gallery
    if not a or not b or a == b or a not in gal.data or b not in gal.data:
        raise HTTPException(status_code=400, detail="Pick two different registered people")
    request.state.dismissed.add(frozenset((a, b)))
    return {"ok": True, **_session_payload(request.state.gallery, request.state.dismissed)}


@api_app.get("/api/session/export.pkl")
async def export_pkl(request: Request):
    if not request.state.gallery.data:
        raise HTTPException(status_code=400, detail="Nothing to export")
    data = request.state.gallery.to_pickle_bytes()
    return RawResponse(
        content=data,
        media_type="application/octet-stream",
        headers={"Content-Disposition": 'attachment; filename="student_embeddings.pkl"'}
    )


@api_app.get("/api/session/export.json")
async def export_json(request: Request):
    return {
        "embeddings": request.state.gallery.to_json(),
        "dismissed": [sorted(p) for p in request.state.dismissed],
    }


@api_app.post("/api/session/import")
async def import_gallery(request: Request):
    body = await request.json()
    try:
        if body.get("json") is not None:
            gal = P.Gallery.from_json(body["json"])
        elif body.get("pkl_b64"):
            gal = P.Gallery.from_pickle_bytes(base64.b64decode(body["pkl_b64"]))
        else:
            raise HTTPException(status_code=400, detail="Send either 'json' or 'pkl_b64'")
    except Exception as exc:
        raise HTTPException(status_code=400, detail=f"Could not load: {exc}")

    dismissed = set()
    for pair in body.get("dismissed") or []:
        if (isinstance(pair, (list, tuple)) and len(pair) == 2
                and all(isinstance(x, str) and x in gal.data for x in pair)
                and pair[0] != pair[1]):
            dismissed.add(frozenset(pair))

    with _SESS_LOCK:
        SESSIONS[request.state.sid]["gallery"] = gal
        SESSIONS[request.state.sid]["dismissed"] = dismissed

    return {"ok": True, **_session_payload(gal, dismissed)}


# --------------------------------------------------------------------------
# Gradio Interface (Mounted at /)
# --------------------------------------------------------------------------
with gr.Blocks(title="Face Recognition Backend API") as demo:
    gr.Markdown(
        """
        # 🚀 Face Recognition Backend API
        **Powered by MTCNN + ArcFace (Retrieval-Augmented Recognition)**
        
        This Hugging Face Space hosts the ML inference API for the attendance application.
        
        - **Status:** Running
        - **Documentation:** Interactive OpenAPI docs available at [`/docs`](/docs)
        - **Frontend:** Connected to your live Vercel web application
        """
    )
    with gr.Row():
        btn = gr.Button("Check Model Readiness")
        out = gr.JSON(label="Readiness Info")
        btn.click(fn=lambda: {**_status, "status": "online"}, outputs=out)

app = gr.mount_gradio_app(api_app, demo, path="/")

if __name__ == "__main__":
    print("[startup] Starting Gradio server", flush=True)

    demo.launch(
        server_name="0.0.0.0",
        server_port=1008,
        show_error=True
    )