import io
import base64
import numpy as np

from models.face_recognition import ArcFaceRecognition, extract_faces



def get_face_predictions(face, arcface_model, data_manager, results):
    """
    Get predictions for a single face
    
    Parameters:
    face: Face image array
    arcface_model: ArcFaceRecognition instance
    data_manager: DataManager instance
    results: Number of results to return
    
    Returns:
    List of (name, confidence) tuples from vector search
    """
    # Create batch with original and flipped images
    face_batch = np.stack([face, face[:, ::-1, :]], axis=0)

    # Get embeddings for both orientations in one batch call
    embeddings_batch = arcface_model.get_face_embeddings_batch(face_batch)
    arc = np.mean(embeddings_batch, axis=0)
    
    # Get predictions from ArcFace model
    query_limit = max(results, 50)
    arc_results = data_manager.query_arc_index(arc, query_limit)
    
    # Convert distances to confidence scores and return top results
    names, distances = arc_results
    confidences = 1.0 - np.array(distances)  # Cosine distance to cosine similarity
    
    # Sort by confidence and return top results
    sorted_indices = np.argsort(confidences)[::-1]
    return [(names[i], confidences[i]) for i in sorted_indices[:results]]


def image_search_performers(image, data_manager, threshold=0.5, results=4):
    """
    Search for multiple performers in an image
    
    Parameters:
    image: PIL Image object
    data_manager: DataManager instance
    threshold: Confidence threshold
    results: Number of results to return
    
    Returns:
    List of dictionaries with face image and performer information
    """
    image_array = np.array(image)
    arcface_model = ArcFaceRecognition()

    try:
        faces = extract_faces(image_array)
    except ValueError:
        raise ValueError("No faces found")

    response = []
    for face in faces:
        predictions = get_face_predictions(face['face'], arcface_model, data_manager, results)
        
        # Crop and encode face image
        area = face['facial_area']
        cimage = image.crop((area['x'], area['y'], area['x'] + area['w'], area['y'] + area['h']))
        buf = io.BytesIO()
        cimage.save(buf, format='JPEG')
        im_b64 = base64.b64encode(buf.getvalue()).decode('ascii')

        # Get performer information
        performers = []
        for name, confidence in predictions:
            performer_info = data_manager.get_performer_info(data_manager.faces[name], confidence)
            if performer_info:
                performers.append(performer_info)

        response.append({
            'image': im_b64,
            'area': area,
            'confidence': face['confidence'],
            'performers': performers
        })
    return response