backend/main.py

import os
import asyncio
import queue
import cv2
import numpy as np
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, status, Depends, Header, Query, UploadFile, File, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from backend.schemas import EmployeeCreate, EnrollmentCaptureRequest, CameraSelectRequest
import backend.database as db
from backend.camera import VideoCamera, attendance_alerts

# local USERS credentials mapping and role scopes
USERS = {
    "admin": "admin123",
    "hr": "hr123",
    "kiosk": "kiosk123"
}

ROLES = {
    "admin": "admin",
    "hr": "hr",
    "kiosk": "kiosk"
}

class LoginRequest(BaseModel):
    username: str
    password: str

def get_current_role(authorization: str = Header(None), token: str = Query(None)):
    """Extracts and validates mock-token from Header or Query parameters."""
    raw_token = None
    if isinstance(authorization, str) and authorization.startswith("Bearer "):
        raw_token = authorization.split(" ")[1]
    elif token:
        raw_token = token
        
    if not raw_token or not raw_token.startswith("mock-token-"):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Session expired or authorization token missing."
        )
        
    # Format: mock-token-{username}-{role}
    parts = raw_token.split("-")
    if len(parts) < 4:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid authorization token format."
        )
        
    username = parts[2]
    role = parts[3]
    
    if username not in USERS or ROLES[username] != role:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Unauthorized security token context."
        )
        
    return role

# Define static directories
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
FRONTEND_DIR = os.path.join(BASE_DIR, "frontend")

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup: Initialize the VideoCamera instance securely
    # (Webcam will start dynamically only upon authorized streaming requests)
    camera = VideoCamera()
    print("[+] SecureAttend Biometric Attendance Node initialized and active.")
    
    yield
    
    # Shutdown: Release webcam resources safely
    camera.stop()
    print("[+] Webcam stream released on app shutdown.")

app = FastAPI(
    title="AI Face Recognition Attendance System",
    description="Production-ready localized attendance module utilizing OpenCV YuNet + SFace",
    version="1.0.0",
    lifespan=lifespan
)

# Configure CORS for local development
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# API Endpoints

@app.post("/api/auth/login")
def auth_login(req: LoginRequest):
    username = req.username.lower().strip()
    password = req.password.strip()
    
    if username in USERS and USERS[username] == password:
        role = ROLES[username]
        return {
            "token": f"mock-token-{username}-{role}",
            "role": role,
            "username": username
        }
        
    raise HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Incorrect username or password. Secure access denied."
    )

# 1. Employees Management
@app.get("/api/employees")
def list_employees(role: str = Depends(get_current_role)):
    return db.get_all_employees()

@app.post("/api/employees", status_code=status.HTTP_201_CREATED)
def create_employee(emp: EmployeeCreate, role: str = Depends(get_current_role)):
    if role not in ["admin", "hr"]:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Access denied. Requires HR or Admin privileges."
        )
    # Check if ID or Email already exists
    existing = db.get_employee(emp.id)
    if existing:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Employee ID already registered."
        )
    
    success = db.add_employee(emp.id, emp.name, emp.email, emp.role)
    if not success:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Failed to register employee. Ensure email is unique."
        )
    return {"message": "Employee registered successfully", "employee_id": emp.id}

@app.delete("/api/employees/{emp_id}")
def remove_employee(emp_id: str, role: str = Depends(get_current_role)):
    if role != "admin":
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Access denied. Requires Admin privileges."
        )
    existing = db.get_employee(emp_id)
    if not existing:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Employee not found."
        )
    
    success = db.delete_employee(emp_id)
    if not success:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Failed to delete employee profile."
        )
        
    # Reload embeddings cache immediately to purge deleted user
    VideoCamera().engine.reload_embeddings()
    return {"message": f"Employee {emp_id} successfully deleted."}

# 2. Multi-Angle Enrollment Capture Endpoint
@app.post("/api/enroll/capture")
def capture_enrollment_frame(req: EnrollmentCaptureRequest, role: str = Depends(get_current_role)):
    if role not in ["admin", "hr"]:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Access denied. Requires HR or Admin privileges."
        )
    camera = VideoCamera()
    
    # 1. Validate employee exists
    emp = db.get_employee(req.employee_id)
    if not emp:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Employee profile must be created first before capturing face scans."
        )
    
    # 2. Pull the latest raw frame and pre-detected faces from the active thread-safe capture loop
    raw_frame, faces = camera.get_latest_face_sample()
    if raw_frame is None:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="Webcam frame is currently loading or unavailable. Make sure webcam is active!"
        )
        
    if faces is None or len(faces) == 0:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail="No face detected in camera view! Position yourself centered in the frame."
        )
        
    if len(faces) > 1:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail="Multiple faces detected. Ensure only the enrolling employee is visible in front of the camera!"
        )
        
    # 4. Extract embedding from the single detected face
    face = faces[0]
    embedding = camera.engine.extract_embedding(raw_frame, face)
    
    if embedding is None:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Failed to extract biometric features. Re-align and try again."
        )
        
    # 5. Save raw bytes in SQLite
    save_success = db.save_face_embedding(req.employee_id, req.angle.lower(), embedding)
    if not save_success:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Failed to write face embedding to database."
        )
        
    # 6. Reload engine's cache in background to integrate new embedding
    camera.engine.reload_embeddings()
    
    # 7. Check if all 5 angles are complete
    return {
        "message": f"Successfully captured {req.angle.upper()} scan!",
        "angle": req.angle,
        "is_fully_enrolled": is_fully_enrolled
    }

@app.post("/api/scan/frame")
async def scan_frame(file: UploadFile = File(...), role: str = Depends(get_current_role)):
    # Read and decode the image file
    contents = await file.read()
    nparr = np.frombuffer(contents, np.uint8)
    frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    
    if frame is None:
        raise HTTPException(status_code=400, detail="Invalid image frame received.")
        
    camera = VideoCamera()
    # Perform face detection
    retval, faces = camera.engine.detect_faces(frame)
    
    results = []
    if faces is not None and len(faces) > 0:
        for face in faces:
            # Parse bounding box and landmarks
            bbox = face[0:4].astype(int).tolist() # [x, y, w, h]
            landmarks = face[4:14].reshape(5, 2).astype(int).tolist() # [[x, y], ...]
            
            # Extract embedding & match
            emp_id, emp_name, match_score = None, "Unknown", 0.0
            query_emb = camera.engine.extract_embedding(frame, face)
            if query_emb is not None:
                emp_id, emp_name, match_score = camera.engine.match_face(query_emb)
            
            # Process attendance check-in/out
            event_status = ""
            if emp_id:
                event_status = camera._process_attendance_trigger(emp_id, emp_name, match_score)
                
            results.append({
                "bbox": bbox,
                "landmarks": landmarks,
                "employee_id": emp_id,
                "name": emp_name,
                "match_score": float(match_score),
                "event_status": event_status
            })
            
    return {"faces": results}

@app.post("/api/enroll/upload-frame")
async def enroll_upload_frame(
    employee_id: str = Form(...),
    angle: str = Form(...),
    file: UploadFile = File(...),
    role: str = Depends(get_current_role)
):
    if role not in ["admin", "hr"]:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Access denied. Requires HR or Admin privileges."
        )
        
    # Decode frame
    contents = await file.read()
    nparr = np.frombuffer(contents, np.uint8)
    frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    
    if frame is None:
        raise HTTPException(status_code=400, detail="Invalid image frame received.")
        
    camera = VideoCamera()
    
    # 1. Validate employee exists
    emp = db.get_employee(employee_id)
    if not emp:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Employee profile must be created first before capturing face scans."
        )
        
    # 2. Detect faces
    retval, faces = camera.engine.detect_faces(frame)
    if faces is None or len(faces) == 0:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail="No face detected in camera view! Position yourself centered in the frame."
        )
    if len(faces) > 1:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail="Multiple faces detected. Ensure only the enrolling employee is visible!"
        )
        
    # 3. Extract embedding
    face = faces[0]
    embedding = camera.engine.extract_embedding(frame, face)
    if embedding is None:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Failed to extract biometric features. Re-align and try again."
        )
        
    # 4. Save embedding
    save_success = db.save_face_embedding(employee_id, angle.lower(), embedding)
    if not save_success:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Failed to write face embedding to database."
        )
        
    # 5. Reload engine's cache
    camera.engine.reload_embeddings()
    
    # 6. Check enrollment status
    is_fully_enrolled = db.has_all_embeddings(employee_id)
    
    return {
        "message": f"Successfully captured {angle.upper()} scan!",
        "angle": angle,
        "is_fully_enrolled": is_fully_enrolled
    }

# 3. Attendance Reports
@app.get("/api/attendance/logs")
def fetch_logs(limit: int = 100, role: str = Depends(get_current_role)):
    return db.get_attendance_history(limit)

@app.get("/api/attendance/stats")
def fetch_stats(role: str = Depends(get_current_role)):
    return db.get_todays_stats()

@app.get("/api/attendance/analytics")
def fetch_analytics(role: str = Depends(get_current_role)):
    return {
        "weekly": db.get_weekly_attendance(),
        "hourly": db.get_hourly_attendance()
    }

# 4. Camera Setup API
@app.post("/api/camera/select")
def select_camera(req: CameraSelectRequest, role: str = Depends(get_current_role)):
    if role != "admin":
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Access denied. Requires Admin privileges."
        )
    camera = VideoCamera()
    success = camera.change_camera(req.index)
    if not success:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Could not connect to camera device index {req.index}."
        )
    return {"message": f"Successfully switched camera to index {req.index}."}

@app.post("/api/camera/toggle")
def toggle_camera(req: dict, role: str = Depends(get_current_role)):
    if role != "admin":
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Access denied. Requires Admin privileges."
        )
    active = req.get("active", True)
    camera = VideoCamera()
    if active:
        success = camera.start()
        if not success:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail="Failed to initialize webcam."
            )
        return {"message": "Camera stream enabled successfully.", "active": True}
    else:
        camera.stop()
        return {"message": "Camera stream disabled successfully.", "active": False}

# 5. Live MJPEG Streaming Route
def mjpeg_generator(camera: VideoCamera):
    while True:
        # Fetch latest processed frame bytes
        frame_bytes = camera.get_processed_frame_bytes()
        # Yield as standard boundary frame
        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')
        # Cap streaming transfer frequency
        time_to_sleep = 0.033
        import time
        time.sleep(time_to_sleep)

@app.get("/api/stream/live")
def get_live_mjpeg_stream(token: str = Query(None)):
    get_current_role(token=token) # Validate token query parameter
    camera = VideoCamera()
    # Force start camera if it's currently stopped
    camera.start()
    return StreamingResponse(
        mjpeg_generator(camera),
        media_type="multipart/x-mixed-replace; boundary=frame"
    )

# 6. Real-Time Alert Event Stream (SSE)
async def alert_event_generator():
    while True:
        try:
            # Poll for any new logs logged by the camera background thread
            # Non-blocking pull from queue
            alert = attendance_alerts.get_nowait()
            import json
            yield f"data: {json.dumps(alert)}\n\n"
        except queue.Empty:
            # Send keep-alive comment every 15 seconds to keep browser connection active
            await asyncio.sleep(0.5)

@app.get("/api/stream/alerts")
def get_alert_stream(token: str = Query(None)):
    get_current_role(token=token) # Validate token query parameter
    return StreamingResponse(
        alert_event_generator(),
        media_type="text/event-stream"
    )

# Mount Frontend Static Assets
if os.path.exists(FRONTEND_DIR):
    app.mount("/static", StaticFiles(directory=FRONTEND_DIR), name="static")
    
    @app.get("/")
    def serve_index():
        return RedirectResponse(url="/static/index.html")
else:
    print(f"[-] Warning: Frontend directory {FRONTEND_DIR} not found. Serve endpoints only.")