"""FaceGuard - Advanced Face Recognition & Attendance Backend

Layers of anti-spoofing:
  1. Client-side MediaPipe Face Mesh liveness challenges (blink, head turn, nod, smile) - defeats static photos
  2. Server-side MiniFASNetV2 CNN anti-spoof texture check - defeats printed photos + screen replays
  3. Server-side InsightFace ArcFace recognition (via uniface) for identity match
"""
import os
import io
import csv
import json
import base64
import asyncio
import secrets
from datetime import datetime, timezone, timedelta
from typing import List, Optional, Iterable, Annotated
from contextlib import asynccontextmanager

import cv2
import numpy as np
import onnxruntime as ort
from bson import ObjectId
from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException, UploadFile, File, Form, Query, Request
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel, Field, BeforeValidator
from starlette.middleware.cors import CORSMiddleware
from starlette.concurrency import run_in_threadpool
from motor.motor_asyncio import AsyncIOMotorClient
from uniface import RetinaFace, ArcFace

load_dotenv()

# ---------------- Config ----------------
MONGO_URL = os.environ["MONGO_URL"]
DB_NAME = os.environ["DB_NAME"]
MODEL_PATH = os.environ.get("MODEL_PATH", "/app/backend/models/2.7_80x80_MiniFASNetV2.onnx")
FACE_SIM_THRESHOLD = float(os.environ.get("FACE_SIM_THRESHOLD", "0.42"))
DETECTION_THRESHOLD = float(os.environ.get("DETECTION_THRESHOLD", "0.7"))
LIVENESS_THRESHOLD = float(os.environ.get("LIVENESS_THRESHOLD", "0.7"))

# ---------------- Model container ----------------
models: dict = {"detector": None, "recognizer": None, "liveness": None}

def _init_models():
    print("[FaceGuard] Loading RetinaFace detector...")
    models["detector"] = RetinaFace()
    print("[FaceGuard] Loading ArcFace recognizer...")
    models["recognizer"] = ArcFace()
    print("[FaceGuard] Loading MiniFASNetV2 liveness (texture anti-spoof)...")
    sess = ort.InferenceSession(MODEL_PATH, providers=["CPUExecutionProvider"])
    models["liveness"] = {
        "session": sess,
        "input_name": sess.get_inputs()[0].name,
    }
    print("[FaceGuard] All models ready.")

# ---------------- Mongo ----------------
mongo_client: Optional[AsyncIOMotorClient] = None
db = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global mongo_client, db
    mongo_client = AsyncIOMotorClient(MONGO_URL)
    db = mongo_client[DB_NAME]
    await db.employees.create_index("employee_id", unique=True)
    await db.attendance.create_index([("employee_id", 1), ("timestamp", -1)])
    try:
        _init_models()
    except Exception as e:
        print(f"[FaceGuard] Model init error: {e}")
    yield
    mongo_client.close()

app = FastAPI(title="FaceGuard API", version="1.0.0", lifespan=lifespan)

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

# ---------------- Utils ----------------
PyObjectId = Annotated[str, BeforeValidator(lambda v: str(v) if isinstance(v, ObjectId) else v)]

def _strip_b64(s: str) -> str:
    if isinstance(s, str) and "," in s and s.lstrip().lower().startswith("data:"):
        return s.split(",", 1)[1]
    return s

def _b64_to_bgr(b64: str) -> Optional[np.ndarray]:
    try:
        raw = base64.b64decode(_strip_b64(b64), validate=False)
        arr = np.frombuffer(raw, np.uint8)
        return cv2.imdecode(arr, cv2.IMREAD_COLOR)
    except Exception:
        return None

def _bgr_to_b64_jpeg(img: np.ndarray, quality: int = 80) -> str:
    ok, buf = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, quality])
    if not ok:
        return ""
    return "data:image/jpeg;base64," + base64.b64encode(buf.tobytes()).decode()

def _softmax(x: np.ndarray) -> np.ndarray:
    e = np.exp(x - np.max(x))
    return e / e.sum()

def _predict_liveness(img_bgr: np.ndarray, bbox) -> tuple:
    """MiniFASNetV2 texture-based anti-spoof.

    Returns (is_real: bool, live_score: float in [0,1])
    """
    if models["liveness"] is None:
        return True, 1.0
    h_img, w_img, _ = img_bgr.shape
    x1, y1, x2, y2 = map(int, bbox)
    box_h = max(1, y2 - y1)
    box_w = max(1, x2 - x1)
    cx = x1 + box_w // 2
    cy = y1 + box_h // 2
    side = int(((box_h + box_w) / 2) * 2.7)
    nx1 = max(0, cx - side // 2)
    ny1 = max(0, cy - side // 2)
    nx2 = min(w_img, cx + side // 2)
    ny2 = min(h_img, cy + side // 2)
    crop = img_bgr[ny1:ny2, nx1:nx2]
    if crop.size == 0:
        return False, 0.0
    resized = cv2.resize(crop, (80, 80)).astype(np.float32)
    blob = np.expand_dims(np.transpose(resized, (2, 0, 1)), 0)
    sess = models["liveness"]["session"]
    inp = models["liveness"]["input_name"]
    out = sess.run(None, {inp: blob})
    probs = _softmax(out[0][0])
    # class index 1 is "real"
    real_score = float(probs[1])
    is_real = bool(np.argmax(probs) == 1 and real_score >= LIVENESS_THRESHOLD)
    return is_real, real_score

def _detect_and_embed(img_bgr: np.ndarray):
    """Detect largest face and compute normalized embedding + spoof check."""
    faces = models["detector"].detect(img_bgr)
    if not faces:
        return None
    # take face with largest area
    def area(f):
        x1, y1, x2, y2 = f.bbox
        return (x2 - x1) * (y2 - y1)
    face = max(faces, key=area)
    if face.confidence < DETECTION_THRESHOLD:
        return {"error": "low_confidence", "det_score": float(face.confidence)}
    emb = models["recognizer"].get_normalized_embedding(img_bgr, face.landmarks)
    if emb is None:
        return {"error": "no_embedding"}
    return {
        "bbox": [float(v) for v in face.bbox],
        "det_score": float(face.confidence),
        "embedding": emb.flatten().astype(np.float32),
        "landmarks": [[float(x), float(y)] for x, y in face.landmarks],
    }

async def _search_match(probe_emb: np.ndarray, employee_id: Optional[str] = None):
    """Search MongoDB embeddings for best cosine similarity match."""
    query = {"employee_id": employee_id} if employee_id else {}
    best_score = -1.0
    best = None
    async for emp in db.employees.find(query):
        for enc in emp.get("embeddings", []):
            ref = np.array(enc, dtype=np.float32)
            sim = float(np.dot(probe_emb, ref))
            if sim > best_score:
                best_score = sim
                best = emp
    return best, best_score

# ---------------- Schemas ----------------
class EmployeeCreate(BaseModel):
    employee_id: str
    name: str
    email: Optional[str] = None
    department: Optional[str] = None
    images: List[str] = Field(..., description="Base64/data-url face images (>=1)")

class EmployeeOut(BaseModel):
    id: PyObjectId = Field(alias="_id")
    employee_id: str
    name: str
    email: Optional[str] = None
    department: Optional[str] = None
    thumbnail: Optional[str] = None
    embeddings_count: int = 0
    created_at: str
    model_config = {"populate_by_name": True, "arbitrary_types_allowed": True}

class VerifyRequest(BaseModel):
    image: str
    employee_id: Optional[str] = None  # if provided, search only this user
    challenges_passed: List[str] = Field(default_factory=list)
    challenge_token: Optional[str] = None

class AttendanceRequest(BaseModel):
    image: str
    type: str  # 'check-in' or 'check-out'
    employee_id: Optional[str] = None
    challenges_passed: List[str] = Field(default_factory=list)
    challenge_token: Optional[str] = None
    latitude: Optional[float] = None
    longitude: Optional[float] = None
    address: Optional[str] = None

# In-memory challenge tokens (short lived)
_challenge_store: dict = {}
CHALLENGE_TTL_SECS = 120

def _issue_challenges():
    import random
    pool = ["blink", "turn_left", "turn_right", "nod", "smile"]
    picked = random.sample(pool, 3)
    tok = secrets.token_urlsafe(18)
    _challenge_store[tok] = {
        "challenges": picked,
        "expires": datetime.now(timezone.utc) + timedelta(seconds=CHALLENGE_TTL_SECS),
        "consumed": False,
    }
    # cleanup old
    now = datetime.now(timezone.utc)
    for k in list(_challenge_store.keys()):
        if _challenge_store[k]["expires"] < now:
            _challenge_store.pop(k, None)
    return tok, picked

def _validate_challenges(token: Optional[str], passed: List[str]) -> tuple:
    if not token or token not in _challenge_store:
        return False, "no_or_invalid_challenge_token"
    entry = _challenge_store[token]
    if entry["consumed"]:
        return False, "challenge_already_used"
    if entry["expires"] < datetime.now(timezone.utc):
        _challenge_store.pop(token, None)
        return False, "challenge_expired"
    required = set(entry["challenges"])
    got = set(passed or [])
    if not required.issubset(got):
        return False, f"missing_challenges:{sorted(required - got)}"
    entry["consumed"] = True
    return True, "ok"

# ---------------- Endpoints ----------------
@app.get("/api/health")
async def health():
    ready = all(models.get(k) is not None for k in ("detector", "recognizer", "liveness"))
    return {"status": "ok" if ready else "loading", "models_ready": ready}

@app.post("/api/attendance/challenge")
async def get_challenge():
    tok, picked = _issue_challenges()
    return {"token": tok, "challenges": picked, "ttl_seconds": CHALLENGE_TTL_SECS}

@app.post("/api/employees")
async def create_employee(payload: EmployeeCreate):
    if not payload.images:
        raise HTTPException(400, "At least one face image required")
    if models["detector"] is None or models["recognizer"] is None:
        raise HTTPException(503, "Models loading, try again")
    embeddings = []
    thumbnail_b64 = None
    for img_str in payload.images:
        img = _b64_to_bgr(img_str)
        if img is None:
            continue
        res = await run_in_threadpool(_detect_and_embed, img)
        if not res or "error" in res:
            continue
        embeddings.append(res["embedding"].tolist())
        if thumbnail_b64 is None:
            # crop face for thumbnail
            x1, y1, x2, y2 = map(int, res["bbox"])
            pad = int(max(x2 - x1, y2 - y1) * 0.2)
            h, w = img.shape[:2]
            cy1 = max(0, y1 - pad)
            cy2 = min(h, y2 + pad)
            cx1 = max(0, x1 - pad)
            cx2 = min(w, x2 + pad)
            face_crop = img[cy1:cy2, cx1:cx2]
            face_crop = cv2.resize(face_crop, (256, 256))
            thumbnail_b64 = _bgr_to_b64_jpeg(face_crop, 82)
    if not embeddings:
        raise HTTPException(400, "No valid face detected in provided images")
    doc = {
        "employee_id": payload.employee_id.strip(),
        "name": payload.name.strip(),
        "email": (payload.email or "").strip() or None,
        "department": (payload.department or "").strip() or None,
        "embeddings": embeddings,
        "embeddings_count": len(embeddings),
        "thumbnail": thumbnail_b64,
        "created_at": datetime.now(timezone.utc).isoformat(),
    }
    try:
        result = await db.employees.insert_one(doc)
    except Exception as e:
        if "duplicate" in str(e).lower():
            raise HTTPException(409, "employee_id already exists")
        raise
    return {
        "id": str(result.inserted_id),
        "employee_id": doc["employee_id"],
        "name": doc["name"],
        "embeddings_count": len(embeddings),
        "thumbnail": thumbnail_b64,
    }

@app.get("/api/employees")
async def list_employees():
    out = []
    async for e in db.employees.find({}, {"embeddings": 0}).sort("created_at", -1):
        out.append({
            "id": str(e["_id"]),
            "employee_id": e["employee_id"],
            "name": e["name"],
            "email": e.get("email"),
            "department": e.get("department"),
            "thumbnail": e.get("thumbnail"),
            "created_at": e.get("created_at"),
            "embeddings_count": e.get("embeddings_count", 0),
        })
    return out

@app.get("/api/employees/{employee_id}")
async def get_employee(employee_id: str):
    e = await db.employees.find_one({"employee_id": employee_id}, {"embeddings": 0})
    if not e:
        raise HTTPException(404, "employee not found")
    e["id"] = str(e.pop("_id"))
    return e

@app.delete("/api/employees/{employee_id}")
async def delete_employee(employee_id: str):
    res = await db.employees.delete_one({"employee_id": employee_id})
    if res.deleted_count == 0:
        raise HTTPException(404, "employee not found")
    return {"deleted": True, "employee_id": employee_id}

@app.post("/api/face/verify")
async def face_verify(payload: VerifyRequest):
    """Verify a single frame -> face match + spoof check.

    challenge_token/challenges_passed enforced when provided."""
    if models["detector"] is None:
        raise HTTPException(503, "Models loading")
    img = _b64_to_bgr(payload.image)
    if img is None:
        raise HTTPException(400, "invalid image")
    if payload.challenge_token:
        ok, reason = _validate_challenges(payload.challenge_token, payload.challenges_passed)
        if not ok:
            return {"authorized": False, "reason": reason, "stage": "liveness_challenge"}
    res = await run_in_threadpool(_detect_and_embed, img)
    if res is None:
        return {"authorized": False, "reason": "no_face_found", "stage": "detection"}
    if "error" in res:
        return {"authorized": False, "reason": res["error"], "stage": "detection", **{k: v for k, v in res.items() if k != "embedding"}}
    # Texture-based anti-spoof
    is_real, live_score = await run_in_threadpool(_predict_liveness, img, res["bbox"])
    if not is_real:
        return {
            "authorized": False,
            "reason": "spoof_detected",
            "stage": "anti_spoof",
            "live_score": round(live_score, 4),
            "bbox": res["bbox"],
        }
    match, score = await _search_match(res["embedding"], payload.employee_id)
    authorized = bool(match and score >= FACE_SIM_THRESHOLD)
    return {
        "authorized": authorized,
        "similarity": round(score, 4),
        "threshold": FACE_SIM_THRESHOLD,
        "live_score": round(live_score, 4),
        "det_score": round(res["det_score"], 4),
        "bbox": res["bbox"],
        "employee": ({
            "employee_id": match["employee_id"],
            "name": match["name"],
            "department": match.get("department"),
            "thumbnail": match.get("thumbnail"),
        } if authorized and match else None),
        "reason": "ok" if authorized else "no_match",
        "stage": "recognition",
    }

@app.post("/api/attendance")
async def record_attendance(payload: AttendanceRequest):
    if payload.type not in ("check-in", "check-out"):
        raise HTTPException(400, "type must be check-in or check-out")
    if models["detector"] is None:
        raise HTTPException(503, "Models loading")
    img = _b64_to_bgr(payload.image)
    if img is None:
        raise HTTPException(400, "invalid image")
    # 1) challenges
    ok, reason = _validate_challenges(payload.challenge_token, payload.challenges_passed)
    if not ok:
        return {"success": False, "reason": reason, "stage": "liveness_challenge"}
    res = await run_in_threadpool(_detect_and_embed, img)
    if res is None:
        return {"success": False, "reason": "no_face_found", "stage": "detection"}
    if "error" in res:
        return {"success": False, "reason": res["error"], "stage": "detection"}
    is_real, live_score = await run_in_threadpool(_predict_liveness, img, res["bbox"])
    if not is_real:
        return {"success": False, "reason": "spoof_detected", "stage": "anti_spoof", "live_score": round(live_score, 4)}
    match, score = await _search_match(res["embedding"], payload.employee_id)
    if not match or score < FACE_SIM_THRESHOLD:
        return {
            "success": False,
            "reason": "no_match",
            "stage": "recognition",
            "similarity": round(score, 4),
            "threshold": FACE_SIM_THRESHOLD,
        }
    # Save attendance
    x1, y1, x2, y2 = map(int, res["bbox"])
    pad = int(max(x2 - x1, y2 - y1) * 0.15)
    h, w = img.shape[:2]
    face_crop = img[max(0, y1 - pad):min(h, y2 + pad), max(0, x1 - pad):min(w, x2 + pad)]
    if face_crop.size > 0:
        face_crop = cv2.resize(face_crop, (240, 240))
    snapshot_b64 = _bgr_to_b64_jpeg(face_crop if face_crop.size > 0 else img, 78)
    now = datetime.now(timezone.utc)
    record = {
        "employee_id": match["employee_id"],
        "employee_name": match["name"],
        "department": match.get("department"),
        "type": payload.type,
        "timestamp": now.isoformat(),
        "date": now.strftime("%Y-%m-%d"),
        "similarity": round(score, 4),
        "live_score": round(live_score, 4),
        "location": {
            "latitude": payload.latitude,
            "longitude": payload.longitude,
            "address": payload.address,
        },
        "snapshot": snapshot_b64,
        "challenges_passed": payload.challenges_passed,
    }
    result = await db.attendance.insert_one(record)
    return {
        "success": True,
        "id": str(result.inserted_id),
        "type": payload.type,
        "timestamp": record["timestamp"],
        "employee": {
            "employee_id": match["employee_id"],
            "name": match["name"],
            "department": match.get("department"),
            "thumbnail": match.get("thumbnail"),
        },
        "similarity": round(score, 4),
        "live_score": round(live_score, 4),
    }

@app.get("/api/attendance")
async def list_attendance(
    date: Optional[str] = Query(None, description="YYYY-MM-DD"),
    employee_id: Optional[str] = None,
    limit: int = Query(200, le=1000),
    include_snapshot: bool = False,
):
    q = {}
    if date:
        q["date"] = date
    if employee_id:
        q["employee_id"] = employee_id
    proj = None if include_snapshot else {"snapshot": 0}
    out = []
    cursor = db.attendance.find(q, proj).sort("timestamp", -1).limit(limit)
    async for r in cursor:
        r["id"] = str(r.pop("_id"))
        out.append(r)
    return out

@app.get("/api/attendance/{att_id}")
async def get_attendance(att_id: str):
    try:
        oid = ObjectId(att_id)
    except Exception:
        raise HTTPException(400, "invalid id")
    r = await db.attendance.find_one({"_id": oid})
    if not r:
        raise HTTPException(404, "not found")
    r["id"] = str(r.pop("_id"))
    return r

@app.get("/api/attendance/stats/summary")
async def attendance_stats(date: Optional[str] = None):
    if not date:
        date = datetime.now(timezone.utc).strftime("%Y-%m-%d")
    total_employees = await db.employees.count_documents({})
    today_records = await db.attendance.find({"date": date}, {"snapshot": 0}).to_list(None)
    # unique check-ins today
    checked_in_ids = set()
    checked_out_ids = set()
    late_ids = set()
    for r in today_records:
        if r["type"] == "check-in":
            checked_in_ids.add(r["employee_id"])
            # Late if check-in after 09:00 UTC (adjust in prod)
            try:
                t = datetime.fromisoformat(r["timestamp"])
                if t.hour >= 9:
                    late_ids.add(r["employee_id"])
            except Exception:
                pass
        elif r["type"] == "check-out":
            checked_out_ids.add(r["employee_id"])
    return {
        "date": date,
        "total_employees": total_employees,
        "present_today": len(checked_in_ids),
        "checked_out": len(checked_out_ids),
        "late_today": len(late_ids),
        "absent_today": max(0, total_employees - len(checked_in_ids)),
        "total_records_today": len(today_records),
    }

@app.get("/api/attendance/export/csv")
async def export_csv(date: Optional[str] = None, employee_id: Optional[str] = None):
    q = {}
    if date:
        q["date"] = date
    if employee_id:
        q["employee_id"] = employee_id
    buf = io.StringIO()
    writer = csv.writer(buf)
    writer.writerow([
        "Timestamp", "Employee ID", "Name", "Department", "Type",
        "Similarity", "Live Score", "Latitude", "Longitude", "Address",
    ])
    async for r in db.attendance.find(q, {"snapshot": 0}).sort("timestamp", -1):
        loc = r.get("location") or {}
        writer.writerow([
            r.get("timestamp", ""),
            r.get("employee_id", ""),
            r.get("employee_name", ""),
            r.get("department", "") or "",
            r.get("type", ""),
            r.get("similarity", ""),
            r.get("live_score", ""),
            loc.get("latitude", "") or "",
            loc.get("longitude", "") or "",
            loc.get("address", "") or "",
        ])
    buf.seek(0)
    return StreamingResponse(
        iter([buf.getvalue()]),
        media_type="text/csv",
        headers={"Content-Disposition": f"attachment; filename=attendance_{date or 'all'}.csv"},
    )