import spaces
from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.responses import JSONResponse
import tempfile
import os
import base64
import io
from PIL import Image
from typing import List, Optional
import uvicorn
from detect_faces import detect_faces

app = FastAPI(
    title="Face Detection API",
    description="API for detecting faces in images using YOLOv8 - Hugging Face Spaces",
    version="1.0.0"
)

def pil_to_base64(pil_image: Image.Image, format: str = "JPEG") -> str:
    """Convert PIL Image to base64 string"""
    buffer = io.BytesIO()
    pil_image.save(buffer, format=format)
    return base64.b64encode(buffer.getvalue()).decode('utf-8')

@spaces.GPU
def process_image_gpu(input_image_path: str, box_margin: int):
    """GPU-accelerated face detection function"""
    return detect_faces(input_image_path, box_margin=box_margin)

@app.post("/detect-faces/")
async def detect_faces_endpoint(
    image: UploadFile = File(..., description="Image file to process"),
    box_margin: int = Form(default=10, description="Box margin for face detection (0-40)"),
    include_confidence: bool = Form(default=True, description="Include confidence scores in response")
):
    """
    Detect faces in an uploaded image
    
    - **image**: Upload an image file (JPEG, PNG, etc.)
    - **box_margin**: Margin around detected faces (0-40)
    - **include_confidence**: Whether to include confidence scores
    
    Returns a list of detected faces as base64-encoded images with confidence scores
    """
    
    # Validate box_margin
    if not 0 <= box_margin <= 40:
        raise HTTPException(status_code=400, detail="box_margin must be between 0 and 40")
    
    # Validate file type (handle None content_type)
    if image.content_type and not image.content_type.startswith('image/'):
        raise HTTPException(status_code=400, detail="File must be an image")
    
    temp_file_path = None
    try:
        # Save uploaded file temporarily
        with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
            content = await image.read()
            temp_file.write(content)
            temp_file_path = temp_file.name
        
        # Process the image using GPU-decorated function
        # detect_faces returns list of tuples: (PIL_Image, confidence_string)
        faces_data = process_image_gpu(temp_file_path, box_margin)
        
        # Convert faces to base64 with confidence scores
        faces_response = []
        for i, (face_image, confidence) in enumerate(faces_data):
            face_data = {
                "id": i,
                "image": pil_to_base64(face_image),
                "format": "base64"
            }
            
            if include_confidence:
                face_data["confidence"] = confidence
                # Extract numeric confidence for sorting
                try:
                    numeric_confidence = float(confidence.replace('%', ''))
                    face_data["confidence_numeric"] = numeric_confidence
                except:
                    face_data["confidence_numeric"] = 0.0
            
            faces_response.append(face_data)
        
        # Sort by confidence if available
        if include_confidence:
            faces_response.sort(key=lambda x: x.get("confidence_numeric", 0), reverse=True)
        
        return JSONResponse(content={
            "status": "success",
            "faces_count": len(faces_response),
            "faces": faces_response,
            "parameters": {
                "box_margin": box_margin,
                "include_confidence": include_confidence
            }
        })
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
    
    finally:
        # Clean up temporary file
        if temp_file_path and os.path.exists(temp_file_path):
            try:
                os.unlink(temp_file_path)
            except:
                pass

@app.post("/extract-faces/")
async def extract_faces_endpoint(
    image: UploadFile = File(..., description="Image file to process"),
    box_margin: int = Form(default=10, description="Box margin for face detection (0-40)"),
    min_confidence: float = Form(default=50.0, description="Minimum confidence threshold (0-100)"),
    return_format: str = Form(default="base64", description="Return format: 'base64' or 'coordinates'")
):
    """
    Extract individual faces from an image (similar to RetinaFace.extract_faces)
    
    - **image**: Upload an image file (JPEG, PNG, etc.)
    - **box_margin**: Margin around detected faces (0-40)
    - **min_confidence**: Minimum confidence threshold for faces to include
    - **return_format**: 'base64' for images, 'coordinates' for bounding boxes
    
    Returns extracted faces above the confidence threshold
    """
    
    # Validate parameters
    if not 0 <= box_margin <= 40:
        raise HTTPException(status_code=400, detail="box_margin must be between 0 and 40")
    
    if not 0 <= min_confidence <= 100:
        raise HTTPException(status_code=400, detail="min_confidence must be between 0 and 100")
    
    if return_format not in ["base64", "coordinates"]:
        raise HTTPException(status_code=400, detail="return_format must be 'base64' or 'coordinates'")
    
    # Validate file type (handle None content_type)
    if image.content_type and not image.content_type.startswith('image/'):
        raise HTTPException(status_code=400, detail="File must be an image")
    
    temp_file_path = None
    try:
        # Save uploaded file temporarily
        with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
            content = await image.read()
            temp_file.write(content)
            temp_file_path = temp_file.name
        
        # Process the image using GPU-decorated function
        faces_data = process_image_gpu(temp_file_path, box_margin)
        
        # Filter faces by confidence and extract
        extracted_faces = []
        valid_faces_count = 0
        
        for i, (face_image, confidence) in enumerate(faces_data):
            try:
                numeric_confidence = float(confidence.replace('%', ''))
                
                # Only include faces above confidence threshold
                if numeric_confidence >= min_confidence:
                    # Check if face is not empty (similar to RetinaFace check)
                    if face_image.size[0] > 0 and face_image.size[1] > 0:
                        valid_faces_count += 1
                        
                        if return_format == "base64":
                            face_data = {
                                "face_id": i,
                                "image": pil_to_base64(face_image),
                                "format": "base64",
                                "confidence": confidence,
                                "confidence_numeric": numeric_confidence,
                                "dimensions": {
                                    "width": face_image.size[0],
                                    "height": face_image.size[1]
                                }
                            }
                        else:  # coordinates format
                            face_data = {
                                "face_id": i,
                                "confidence": confidence,
                                "confidence_numeric": numeric_confidence,
                                "dimensions": {
                                    "width": face_image.size[0],
                                    "height": face_image.size[1]
                                }
                            }
                        
                        extracted_faces.append(face_data)
                    else:
                        print(f"Skipping empty face {i}")  # Similar to your original code
                        
            except ValueError:
                continue  # Skip faces with invalid confidence values
        
        # Sort by confidence (highest first)
        extracted_faces.sort(key=lambda x: x["confidence_numeric"], reverse=True)
        
        return JSONResponse(content={
            "status": "success",
            "total_faces_detected": len(faces_data),
            "valid_faces_extracted": valid_faces_count,
            "faces_above_threshold": len(extracted_faces),
            "extracted_faces": extracted_faces,
            "parameters": {
                "box_margin": box_margin,
                "min_confidence_threshold": min_confidence,
                "return_format": return_format
            }
        })
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error extracting faces: {str(e)}")
    
    finally:
        # Clean up temporary file
        if temp_file_path and os.path.exists(temp_file_path):
            try:
                os.unlink(temp_file_path)
            except:
                pass

@app.post("/detect-faces-summary/")
async def detect_faces_summary_endpoint(
    image: UploadFile = File(..., description="Image file to process"),
    box_margin: int = Form(default=10, description="Box margin for face detection (0-40)"),
    min_confidence: float = Form(default=50.0, description="Minimum confidence threshold (0-100)")
):
    """
    Detect faces and return only summary information (no images)
    
    - **image**: Upload an image file (JPEG, PNG, etc.)
    - **box_margin**: Margin around detected faces (0-40)
    - **min_confidence**: Minimum confidence threshold for faces to include
    
    Returns summary statistics about detected faces
    """
    
    # Validate parameters
    if not 0 <= box_margin <= 40:
        raise HTTPException(status_code=400, detail="box_margin must be between 0 and 40")
    
    if not 0 <= min_confidence <= 100:
        raise HTTPException(status_code=400, detail="min_confidence must be between 0 and 100")
    
    # Validate file type (handle None content_type)
    if image.content_type and not image.content_type.startswith('image/'):
        raise HTTPException(status_code=400, detail="File must be an image")
    
    temp_file_path = None
    try:
        # Save uploaded file temporarily
        with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
            content = await image.read()
            temp_file.write(content)
            temp_file_path = temp_file.name
        
        # Process the image using GPU-decorated function
        faces_data = process_image_gpu(temp_file_path, box_margin)
        
        # Process confidence scores
        confidences = []
        high_confidence_faces = 0
        
        for face_image, confidence in faces_data:
            try:
                numeric_confidence = float(confidence.replace('%', ''))
                confidences.append(numeric_confidence)
                if numeric_confidence >= min_confidence:
                    high_confidence_faces += 1
            except:
                confidences.append(0.0)
        
        # Calculate statistics
        total_faces = len(faces_data)
        avg_confidence = sum(confidences) / len(confidences) if confidences else 0
        max_confidence = max(confidences) if confidences else 0
        min_confidence_found = min(confidences) if confidences else 0
        
        return JSONResponse(content={
            "status": "success",
            "summary": {
                "total_faces": total_faces,
                "high_confidence_faces": high_confidence_faces,
                "average_confidence": round(avg_confidence, 2),
                "max_confidence": max_confidence,
                "min_confidence": min_confidence_found,
                "confidence_threshold": min_confidence
            },
            "parameters": {
                "box_margin": box_margin,
                "min_confidence_threshold": min_confidence
            }
        })
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
    
    finally:
        # Clean up temporary file
        if temp_file_path and os.path.exists(temp_file_path):
            try:
                os.unlink(temp_file_path)
            except:
                pass

@app.get("/")
async def root():
    """Root endpoint with API information"""
    return {
        "message": "YOLOv8 Face Detection API - Hugging Face Spaces",
        "version": "1.0.0",
        "platform": "Hugging Face Spaces",
        "model": "YOLOv8l-face",
        "endpoints": {
            "/detect-faces/": "POST - Detect faces and return base64 images with confidence",
            "/extract-faces/": "POST - Extract individual faces (RetinaFace alternative)",
            "/detect-faces-summary/": "POST - Detect faces and return summary statistics only",
            "/docs": "GET - Interactive API documentation",
            "/health": "GET - Health check"
        },
        "usage": {
            "supported_formats": ["JPEG", "PNG", "BMP", "TIFF"],
            "max_box_margin": 40,
            "confidence_range": "0-100%"
        }
    }

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    import torch
    return {
        "status": "healthy",
        "platform": "Hugging Face Spaces",
        "gpu_available": torch.cuda.is_available(),
        "cuda_device_count": torch.cuda.device_count() if torch.cuda.is_available() else 0,
        "model_loaded": True
    }

# For Hugging Face Spaces deployment
if __name__ == "__main__":
    # Use the port that Spaces expects (usually 7860)
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port)