"""
arcface_recognizer.py -- ArcFace face embedding extraction and gallery matching.

Performs face alignment using 5-point landmarks (insightface standard reference),
extracts 512-dimensional embeddings via ArcFace (TRT FP16 or ONNX fallback),
and matches against a persistent gallery using cosine similarity.
"""

import os
import logging
from typing import Optional

import numpy as np
import cv2

logger = logging.getLogger(__name__)

# InsightFace standard reference landmarks for 112x112 alignment
ARCFACE_SRC = np.array([
    [38.2946, 51.6963],   # left eye
    [73.5318, 51.5014],   # right eye
    [56.0252, 71.7366],   # nose
    [41.5493, 92.3655],   # left mouth
    [70.7299, 92.2041],   # right mouth
], dtype=np.float32)


# -- Inference backends --------------------------------------------------------

class _TRTBackend:
    """TensorRT inference engine for ArcFace."""

    def __init__(self, engine_path: str):
        import tensorrt as trt
        import pycuda.driver as cuda
        import pycuda.autoinit  # noqa: F401

        self._cuda = cuda
        trt_logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, 'rb') as f, trt.Runtime(trt_logger) as runtime:
            self._engine = runtime.deserialize_cuda_engine(f.read())
        self._context = self._engine.create_execution_context()

        self._inputs = []
        self._outputs = []
        self._bindings = []
        for i in range(self._engine.num_io_tensors):
            name = self._engine.get_tensor_name(i)
            dtype = trt.nptype(self._engine.get_tensor_dtype(name))
            shape = tuple(self._engine.get_tensor_shape(name))
            nbytes = int(np.prod(shape)) * np.dtype(dtype).itemsize
            host_mem = cuda.pagelocked_empty(shape, dtype)
            device_mem = cuda.mem_alloc(nbytes)
            self._bindings.append(int(device_mem))
            if self._engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                self._inputs.append({'host': host_mem, 'device': device_mem})
            else:
                self._outputs.append({'host': host_mem, 'device': device_mem,
                                      'shape': shape})
        self._stream = cuda.Stream()

    def infer(self, input_data: np.ndarray) -> np.ndarray:
        """Run inference and return the embedding vector."""
        np.copyto(self._inputs[0]['host'], input_data.ravel())
        self._cuda.memcpy_htod_async(
            self._inputs[0]['device'], self._inputs[0]['host'], self._stream)
        self._context.execute_async_v2(self._bindings, self._stream.handle)
        for out in self._outputs:
            self._cuda.memcpy_dtoh_async(out['host'], out['device'], self._stream)
        self._stream.synchronize()
        return self._outputs[0]['host'].reshape(self._outputs[0]['shape']).copy()


class _ONNXBackend:
    """ONNX Runtime inference (CUDA EP with CPU fallback)."""

    def __init__(self, onnx_path: str):
        import onnxruntime as ort
        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
        self._session = ort.InferenceSession(onnx_path, providers=providers)
        self._input_name = self._session.get_inputs()[0].name

    def infer(self, input_data: np.ndarray) -> np.ndarray:
        """Run inference and return the embedding vector."""
        results = self._session.run(None, {self._input_name: input_data})
        return results[0]


# -- Face alignment ------------------------------------------------------------

def align_face(bgr: np.ndarray, landmarks_10: list[float]) -> np.ndarray:
    """Align a face to 112x112 using 5-point landmarks.

    Args:
        bgr: Source BGR image.
        landmarks_10: Flat list of 10 floats [x0,y0, x1,y1, ..., x4,y4].

    Returns:
        Aligned BGR face crop of shape (112, 112, 3).
    """
    src_pts = np.array(landmarks_10, dtype=np.float32).reshape(5, 2)
    M, _ = cv2.estimateAffinePartial2D(src_pts, ARCFACE_SRC)
    if M is None:
        # Fallback: simple crop and resize from bbox-like region
        cx = np.mean(src_pts[:, 0])
        cy = np.mean(src_pts[:, 1])
        spread = max(np.ptp(src_pts[:, 0]), np.ptp(src_pts[:, 1])) * 1.5
        half = spread / 2
        x1 = max(0, int(cx - half))
        y1 = max(0, int(cy - half))
        x2 = min(bgr.shape[1], int(cx + half))
        y2 = min(bgr.shape[0], int(cy + half))
        crop = bgr[y1:y2, x1:x2]
        return cv2.resize(crop, (112, 112), interpolation=cv2.INTER_LINEAR)

    aligned = cv2.warpAffine(bgr, M, (112, 112), borderMode=cv2.BORDER_REPLICATE)
    return aligned


# -- Main recognizer class -----------------------------------------------------

class ArcFaceRecognizer:
    """ArcFace face embedding extractor and gallery matcher.

    Args:
        engine_path: Path to TensorRT engine file.
        onnx_path: Path to ONNX model file (used if engine not available).
    """

    def __init__(self, engine_path: str = '', onnx_path: str = ''):
        self._backend: Optional[_TRTBackend | _ONNXBackend] = None
        self.gallery: dict[int, dict] = {}

        # Try TRT first, then ONNX
        if engine_path and os.path.isfile(engine_path):
            try:
                self._backend = _TRTBackend(engine_path)
                logger.info('ArcFace TensorRT backend loaded: %s', engine_path)
                return
            except Exception as e:
                logger.warning('ArcFace TRT load failed (%s), trying ONNX', e)

        if onnx_path and os.path.isfile(onnx_path):
            try:
                self._backend = _ONNXBackend(onnx_path)
                logger.info('ArcFace ONNX backend loaded: %s', onnx_path)
                return
            except Exception as e:
                logger.error('ArcFace ONNX load failed: %s', e)

        logger.error('No ArcFace model loaded. Recognition will be unavailable.')

    @property
    def is_loaded(self) -> bool:
        """Return True if a backend is loaded and ready."""
        return self._backend is not None

    def embed(self, bgr_face_112x112: np.ndarray) -> np.ndarray:
        """Extract 512-dim L2-normalized embedding from a 112x112 aligned face.

        Args:
            bgr_face_112x112: Aligned face crop, BGR, shape (112, 112, 3).

        Returns:
            L2-normalized embedding of shape (512,).
        """
        if self._backend is None:
            return np.zeros(512, dtype=np.float32)

        # Preprocess: BGR->RGB, /255, subtract 0.5, divide 0.5 -> [1,3,112,112]
        rgb = cv2.cvtColor(bgr_face_112x112, cv2.COLOR_BGR2RGB).astype(np.float32)
        rgb = rgb / 255.0
        rgb = (rgb - 0.5) / 0.5
        blob = rgb.transpose(2, 0, 1)[np.newaxis]  # [1, 3, 112, 112]
        blob = np.ascontiguousarray(blob)

        output = self._backend.infer(blob)
        embedding = output.flatten()[:512].astype(np.float32)

        # L2 normalize
        norm = np.linalg.norm(embedding)
        if norm > 0:
            embedding = embedding / norm
        return embedding

    def align_and_embed(self, bgr_image: np.ndarray, landmarks_10: list[float]) -> np.ndarray:
        """Align face using landmarks and extract embedding.

        Args:
            bgr_image: Full BGR image.
            landmarks_10: Flat list of 10 floats from SCRFD detection.

        Returns:
            L2-normalized embedding of shape (512,).
        """
        aligned = align_face(bgr_image, landmarks_10)
        return self.embed(aligned)

    def load_gallery(self, gallery_path: str) -> None:
        """Load gallery from .npz file with JSON metadata sidecar.

        Args:
            gallery_path: Path to the .npz gallery file.
        """
        import json

        if not os.path.isfile(gallery_path):
            logger.info('No gallery file at %s, starting empty.', gallery_path)
            self.gallery = {}
            return

        data = np.load(gallery_path, allow_pickle=False)
        meta_path = gallery_path.replace('.npz', '_meta.json')

        if os.path.isfile(meta_path):
            with open(meta_path, 'r') as f:
                meta = json.load(f)
        else:
            meta = {}

        self.gallery = {}
        for key in data.files:
            pid = int(key)
            embedding = data[key].astype(np.float32)
            norm = np.linalg.norm(embedding)
            if norm > 0:
                embedding = embedding / norm
            info = meta.get(str(pid), {})
            self.gallery[pid] = {
                'name': info.get('name', f'person_{pid}'),
                'embedding': embedding,
                'samples': info.get('samples', 1),
                'enrolled_at': info.get('enrolled_at', 0.0),
            }
        logger.info('Gallery loaded: %d persons from %s', len(self.gallery), gallery_path)

    def save_gallery(self, gallery_path: str) -> None:
        """Save gallery to .npz file with JSON metadata sidecar.

        Args:
            gallery_path: Path to the .npz gallery file.
        """
        import json

        arrays = {}
        meta = {}
        for pid, info in self.gallery.items():
            arrays[str(pid)] = info['embedding']
            meta[str(pid)] = {
                'name': info['name'],
                'samples': info['samples'],
                'enrolled_at': info['enrolled_at'],
            }

        os.makedirs(os.path.dirname(gallery_path) or '.', exist_ok=True)
        np.savez(gallery_path, **arrays)

        meta_path = gallery_path.replace('.npz', '_meta.json')
        with open(meta_path, 'w') as f:
            json.dump(meta, f, indent=2)
        logger.info('Gallery saved: %d persons to %s', len(self.gallery), gallery_path)

    def match(self, embedding: np.ndarray, threshold: float = 0.35) -> tuple[int, str, float]:
        """Match an embedding against the gallery.

        Args:
            embedding: L2-normalized query embedding of shape (512,).
            threshold: Minimum cosine similarity for a match.

        Returns:
            (person_id, person_name, score) or (-1, '', 0.0) if no match.
        """
        if not self.gallery:
            return (-1, '', 0.0)

        best_pid = -1
        best_name = ''
        best_score = 0.0

        for pid, info in self.gallery.items():
            score = float(np.dot(embedding, info['embedding']))
            if score > best_score:
                best_score = score
                best_pid = pid
                best_name = info['name']

        if best_score >= threshold:
            return (best_pid, best_name, best_score)
        return (-1, '', 0.0)

    def enroll(self, person_id: int, person_name: str, embeddings_list: list[np.ndarray]) -> None:
        """Enroll a person by averaging multiple embeddings.

        Args:
            person_id: Unique integer ID for this person.
            person_name: Human-readable name.
            embeddings_list: List of L2-normalized embeddings to average.
        """
        import time as _time

        if not embeddings_list:
            return

        mean_emb = np.mean(embeddings_list, axis=0).astype(np.float32)
        norm = np.linalg.norm(mean_emb)
        if norm > 0:
            mean_emb = mean_emb / norm

        self.gallery[person_id] = {
            'name': person_name,
            'embedding': mean_emb,
            'samples': len(embeddings_list),
            'enrolled_at': _time.time(),
        }
        logger.info('Enrolled person %d (%s) with %d samples.',
                     person_id, person_name, len(embeddings_list))
