# app.py - FaceNet/ArcFace Face Verification API
from fastapi import FastAPI, File, UploadFile, HTTPException, Body, Form
from fastapi.responses import JSONResponse
import torch
import numpy as np
import cv2
from PIL import Image
import tempfile
import os
import logging
from typing import List
import json
from facenet_pytorch import InceptionResnetV1, extract_face
from facenet_pytorch import MTCNN
import torch.nn.functional as F

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="Face Verification API",
    description="Extract face features from videos using FaceNet model - Specialized for Face Verification",
    version="2.0.0"
)

# Global variables for model
mtcnn = None
resnet = None
device = None

@app.on_event("startup")
async def load_model():
    """Load the FaceNet model when the application starts"""
    global mtcnn, resnet, device
    
    logger.info("Loading FaceNet model and MTCNN face detector...")
    try:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        logger.info(f"Using device: {device}")
        
        # MTCNN for face detection (much better than Haar Cascade)
        mtcnn = MTCNN(image_size=160, margin=0, min_face_size=20, device=device)
        logger.info("✅ MTCNN face detector loaded!")
        
        # FaceNet model (InceptionResnetV1) - trained on VGGFace2
        # Returns 512-dimensional embeddings
        resnet = InceptionResnetV1(pretrained='vggface2', device=device).eval()
        logger.info("✅ FaceNet (InceptionResnetV1) model loaded!")
        logger.info("Model trained on VGGFace2 - specifically for face verification")
        
    except Exception as e:
        logger.error(f"❌ Error loading model: {str(e)}")
        raise e

def extract_faces_from_video(video_path: str, frame_interval: int = 5, max_faces: int = None):
    """Extract and align face images from video frames using MTCNN"""
    try:
        cap = cv2.VideoCapture(video_path)
        faces = []
        frame_idx = 0
        total_frames = 0
        faces_detected = 0

        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            if frame_idx % frame_interval == 0:
                # Convert BGR to RGB
                rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                pil_image = Image.fromarray(rgb_frame)
                
                try:
                    # MTCNN detects and aligns faces
                    face_tensor = mtcnn(pil_image)
                    
                    if face_tensor is not None:
                        # face_tensor shape: (1, 3, 160, 160) or (batch, 3, 160, 160)
                        if face_tensor.dim() == 3:
                            face_tensor = face_tensor.unsqueeze(0)
                        
                        faces.append(face_tensor)
                        faces_detected += 1
                        
                except Exception as e:
                    logger.warning(f"Could not detect face in frame {frame_idx}: {str(e)}")
                
                total_frames += 1
            
            frame_idx += 1

        cap.release()
        
        logger.info(f"Video processing: Processed {total_frames} frames, detected {faces_detected} faces")
        
        if len(faces) == 0:
            logger.warning("No faces detected in video")
            return None
        
        # Limit number of faces if specified
        if max_faces and len(faces) > max_faces:
            faces = faces[:max_faces]
        
        # Concatenate all face tensors
        faces_batch = torch.cat(faces, dim=0)
        return faces_batch
    
    except Exception as e:
        logger.error(f"Error processing video: {str(e)}")
        return None

def extract_facenet_embeddings(faces_batch):
    """Extract FaceNet embeddings (512-dim) from face batch"""
    try:
        if faces_batch is None:
            return None
        
        # Move to device
        faces_batch = faces_batch.to(device)
        
        with torch.no_grad():
            # Get 512-dimensional embeddings from FaceNet
            embeddings = resnet(faces_batch)  # Shape: (batch_size, 512)
            
            # L2 normalize embeddings
            embeddings = F.normalize(embeddings, p=2, dim=1)
            
            # Average embeddings from all frames
            avg_embedding = embeddings.mean(dim=0).cpu().numpy()
            
            logger.info(f"Extracted {embeddings.shape[0]} face embeddings, averaging to get final embedding")
        
        return avg_embedding
    
    except Exception as e:
        logger.error(f"Error extracting embeddings: {str(e)}")
        return None

def calculate_cosine_similarity(embedding1: np.ndarray, embedding2: np.ndarray) -> float:
    """Calculate cosine similarity between two normalized embeddings"""
    # Embeddings should already be normalized by FaceNet
    similarity = np.dot(embedding1, embedding2)
    return float(np.clip(similarity, -1.0, 1.0))

@app.get("/")
async def root():
    """Root endpoint with API information"""
    return {
        "message": "🚀 Face Verification API is running!",
        "status": "healthy",
        "model": "FaceNet (InceptionResnetV1) trained on VGGFace2",
        "embedding_dimension": 512,
        "accuracy_improvement": "99.65% on LFW (using face-specialized model)",
        "endpoints": {
            "health": "/health",
            "extract_features": "/extract",
            "verify_with_embedding": "/verify-with-embedding",
            "verify_two_videos": "/verify-two-videos"
        },
        "recommended_threshold": 0.6,
        "usage": "Use POST requests to extract or verify endpoints"
    }

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "model_loaded": resnet is not None,
        "device": str(device),
        "model_type": "FaceNet (InceptionResnetV1)",
        "embedding_size": 512,
        "timestamp": np.datetime64('now').astype(str)
    }

@app.post("/extract")
async def extract_features(file: UploadFile = File(...)):
    """
    Extract facial features from a video file using FaceNet
    
    - **file**: Video file (mp4, avi, mov, mkv)
    - **returns**: Embedding vector of length 512 (FaceNet embeddings)
    """
    try:
        # Validate file type
        if not file.content_type.startswith('video/') and not file.filename.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
            raise HTTPException(
                status_code=400, 
                detail="Only video files are supported (mp4, avi, mov, mkv)"
            )

        # Create temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp:
            content = await file.read()
            if len(content) == 0:
                raise HTTPException(status_code=400, detail="Empty file")
            tmp.write(content)
            tmp_path = tmp.name

        try:
            # Extract faces from video
            faces_batch = extract_faces_from_video(tmp_path)
            
            if faces_batch is None:
                raise HTTPException(
                    status_code=400, 
                    detail="Could not process video or no faces detected. Ensure video contains clear face(s)."
                )
            
            # Extract FaceNet embeddings
            emb = extract_facenet_embeddings(faces_batch)

            if emb is None:
                raise HTTPException(
                    status_code=400, 
                    detail="Could not extract face embeddings"
                )

            return {
                "success": True,
                "embedding": emb.tolist(), 
                "embedding_length": len(emb),
                "embedding_model": "FaceNet (InceptionResnetV1)",
                "message": "Features extracted successfully. Store this embedding in your database."
            }
        finally:
            # Cleanup
            if os.path.exists(tmp_path):
                os.remove(tmp_path)

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Extract features error: {str(e)}")
        raise HTTPException(
            status_code=500, 
            detail=f"Internal server error: {str(e)}"
        )

@app.post("/verify-with-embedding")
async def verify_with_embedding(
    file: UploadFile = File(..., description="New video file to verify"),
    stored_embedding: str = Form(..., description="Stored embedding as JSON string"),
    threshold: float = Form(0.6, description="Similarity threshold (0.0-1.0), default 0.6")
):
    """
    Verify a new video against a stored FaceNet embedding from database
    
    - **file**: New video file to verify (FormData)
    - **stored_embedding**: Pre-existing embedding array from database as JSON string (FormData)
    - **threshold**: Similarity threshold (default: 0.6 for FaceNet) (FormData)
    - **returns**: Similarity score and verification result
    """
    try:
        # Validate file type
        if not file.content_type.startswith('video/') and not file.filename.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
            raise HTTPException(
                status_code=400, 
                detail="Only video files are supported (mp4, avi, mov, mkv)"
            )

        # Parse stored_embedding from JSON string
        try:
            stored_emb_list = json.loads(stored_embedding)
            stored_emb = np.array(stored_emb_list, dtype=np.float32)
        except Exception as e:
            raise HTTPException(
                status_code=400,
                detail=f"Invalid embedding format: {str(e)}"
            )

        # Validate stored embedding (FaceNet uses 512 dimensions)
        if not stored_emb_list or len(stored_emb_list) != 512:
            raise HTTPException(
                status_code=400,
                detail="Stored embedding must be a 512-dimensional FaceNet embedding"
            )

        # Validate threshold
        if not 0.0 <= threshold <= 1.0:
            raise HTTPException(
                status_code=400,
                detail="Threshold must be between 0.0 and 1.0"
            )

        # Create temporary file for new video
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp:
            content = await file.read()
            if len(content) == 0:
                raise HTTPException(status_code=400, detail="Video file is empty")
            tmp.write(content)
            tmp_path = tmp.name

        try:
            # Extract faces from new video
            logger.info(f"Processing video file: {file.filename}")
            faces_batch = extract_faces_from_video(tmp_path)

            if faces_batch is None:
                raise HTTPException(
                    status_code=400, 
                    detail="Could not process the video or no faces detected"
                )

            # Extract FaceNet embeddings
            new_emb = extract_facenet_embeddings(faces_batch)
            
            if new_emb is None:
                raise HTTPException(
                    status_code=400, 
                    detail="Could not extract face embeddings"
                )

            # Calculate cosine similarity
            similarity = calculate_cosine_similarity(stored_emb, new_emb)
            
            # Determine result
            result = "same_person" if similarity >= threshold else "different_person"

            logger.info(f"Verification completed - Similarity: {similarity:.4f}, Threshold: {threshold}, Result: {result}")

            return {
                "success": True,
                "similarity": round(similarity, 4),
                "threshold_used": threshold,
                "result": result,
                "is_match": similarity >= threshold,
                "model": "FaceNet (InceptionResnetV1) on VGGFace2",
                "message": "Verification completed successfully",
                "comparison": {
                    "stored_embedding_length": len(stored_emb),
                    "new_embedding_length": len(new_emb),
                    "file_processed": file.filename
                }
            }

        finally:
            # Cleanup temp file
            if os.path.exists(tmp_path):
                os.remove(tmp_path)

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Verify with embedding error: {str(e)}")
        raise HTTPException(
            status_code=500, 
            detail=f"Internal server error: {str(e)}"
        )

@app.post("/verify-two-videos")
async def verify_two_videos(
    file1: UploadFile = File(..., description="First video file"),
    file2: UploadFile = File(..., description="Second video file"),
    threshold: float = Body(0.6, description="Similarity threshold (0.0-1.0), default 0.6")
):
    """
    Verify if two videos contain the same person using FaceNet
    
    - **file1**: First video file
    - **file2**: Second video file  
    - **threshold**: Similarity threshold (default: 0.6 for FaceNet)
    - **returns**: Similarity score and verification result
    """
    try:
        # Validate files
        for file in [file1, file2]:
            if not file.content_type.startswith('video/') and not file.filename.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
                raise HTTPException(
                    status_code=400, 
                    detail="Only video files are supported (mp4, avi, mov, mkv)"
                )

        # Validate threshold
        if not 0.0 <= threshold <= 1.0:
            raise HTTPException(
                status_code=400,
                detail="Threshold must be between 0.0 and 1.0"
            )

        temp_files = []
        try:
            # Save files temporarily
            for i, file in enumerate([file1, file2]):
                with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp:
                    content = await file.read()
                    if len(content) == 0:
                        raise HTTPException(status_code=400, detail=f"File {i+1} is empty")
                    tmp.write(content)
                    temp_files.append(tmp.name)

            # Extract faces from both videos
            faces_batch1 = extract_faces_from_video(temp_files[0])
            faces_batch2 = extract_faces_from_video(temp_files[1])

            if faces_batch1 is None or faces_batch2 is None:
                raise HTTPException(
                    status_code=400, 
                    detail="Could not process one or both videos or no faces detected"
                )

            # Extract embeddings
            emb1 = extract_facenet_embeddings(faces_batch1)
            emb2 = extract_facenet_embeddings(faces_batch2)
            
            if emb1 is None or emb2 is None:
                raise HTTPException(
                    status_code=400, 
                    detail="Could not extract embeddings from one or both videos"
                )

            # Calculate cosine similarity
            similarity = calculate_cosine_similarity(emb1, emb2)
            
            # Determine result
            result = "same_person" if similarity >= threshold else "different_person"

            logger.info(f"Two-video verification - Similarity: {similarity:.4f}, Result: {result}")

            return {
                "success": True,
                "similarity": round(similarity, 4),
                "threshold_used": threshold,
                "result": result,
                "is_match": similarity >= threshold,
                "model": "FaceNet (InceptionResnetV1) on VGGFace2",
                "message": "Verification completed successfully"
            }

        finally:
            # Cleanup temp files
            for temp_file in temp_files:
                if os.path.exists(temp_file):
                    os.remove(temp_file)

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Verify two videos error: {str(e)}")
        raise HTTPException(
            status_code=500, 
            detail=f"Internal server error: {str(e)}"
        )

# Add CORS middleware
from fastapi.middleware.cors import CORSMiddleware

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