import io
import base64
import numpy as np
from PIL import Image

from models.face_recognition import EnsembleFaceRecognition, extract_faces


def get_face_predictions(face, ensemble, data_manager, results):
    """
    Get predictions for a single face

    Parameters:
    face: Face image array
    ensemble: EnsembleFaceRecognition instance
    data_manager: DataManager instance
    results: Number of results to return

    Returns:
    List of (name, confidence, distance) tuples
    """
    face_batch = np.stack([face, face[:, ::-1, :]], axis=0)
    embeddings_batch = ensemble.get_face_embeddings_batch(face_batch)
    facenet = np.mean(embeddings_batch['facenet'], axis=0)
    arc = np.mean(embeddings_batch['arc'], axis=0)

    ids, scores = data_manager.query_multi(facenet, arc, max(results, 50))
    if not ids:
        return []

    # WeightedReRanker normalizes each field to [0,1] then sums; with two
    # equal-weight fields the max combined score is 2.0. Dividing by 2.0
    # gives an absolute average similarity, and the sqrt curve matches the
    # old distance_to_score shape.
    return [(name, (score / 2.0) ** 0.5, 0.0) for name, score in zip(ids, scores)]


def image_search_performers(image, data_manager, threshold=0.5, results=3):
    """
    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)
    ensemble = EnsembleFaceRecognition({"facenet": 0.5, "arc": 0.5})

    faces = extract_faces(image_array)
    if not faces:
        raise ValueError("No faces found")

    response = []
    for face in faces:
        predictions = get_face_predictions(face['face'], ensemble, data_manager, results)

        cimage = Image.fromarray(face['face'])
        buf = io.BytesIO()
        cimage.save(buf, format='JPEG')
        im_b64 = base64.b64encode(buf.getvalue()).decode('ascii')

        performers = []
        for name, confidence, distance in predictions:
            if confidence < threshold:
                break
            if len(performers) >= results:
                break
            performer_info = data_manager.get_performer_info(name, confidence, distance)
            if performer_info:
                performers.append(performer_info)

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