import os
import numpy as np
import cv2
import torch
from PIL import Image

# Make sure these modules exist in the same package folder
from .exp_recognition_model import trnscm, load_model, classes

#############################################################################################################################
#   Caution: Don't change any of the filenames, function names and definitions                                              #
#   Always use the current_path + file_name for refering any files, without it we cannot access files on the server         #
#############################################################################################################################

current_path = os.path.dirname(os.path.abspath(__file__))

_MODEL = None
_DEVICE = None


def _to_bgr_np(img):
    """
    Convert various inputs -> OpenCV BGR numpy array.
    Supports: filepath str, PIL.Image, numpy RGB/gray/RGBA.
    """
    if img is None:
        return None

    # filepath string
    if isinstance(img, str):
        return cv2.imread(img)

    # PIL image
    if isinstance(img, Image.Image):
        rgb = np.array(img.convert("RGB"))
        return cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)

    # numpy array
    if isinstance(img, np.ndarray):
        arr = img

        # grayscale (H,W)
        if arr.ndim == 2:
            return cv2.cvtColor(arr, cv2.COLOR_GRAY2BGR)

        # RGBA -> RGB
        if arr.ndim == 3 and arr.shape[2] == 4:
            arr = arr[:, :, :3]

        # assume RGB (common from Gradio type="numpy")
        if arr.ndim == 3 and arr.shape[2] == 3:
            return cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)

        return arr

    return None


def detected_face(image_bgr):
    """
    Returns PIL grayscale cropped face with maximum area.
    Returns 0 if not detected or cascade not available.
    """
    face_haar = os.path.join(current_path, "haarcascade_frontalface_default.xml")
    face_cascade = cv2.CascadeClassifier(face_haar)

    # if cascade missing or failed to load
    if face_cascade.empty():
        return 0

    gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray, 1.3, 5)

    if faces is None or len(faces) == 0:
        return 0

    # pick max area
    x, y, w, h = max(faces, key=lambda f: f[2] * f[3])
    crop = gray[y:y + h, x:x + w]
    return Image.fromarray(crop)


def _load_cached_model():
    """
    Load the expression model once and cache it.
    """
    global _MODEL, _DEVICE
    if _MODEL is not None:
        return _MODEL, _DEVICE

    _DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    ckpt_path = os.path.join(current_path, "best_resnet18_expression.pt")
    if not os.path.exists(ckpt_path):
        raise FileNotFoundError(f"Missing checkpoint: {ckpt_path}")

    _MODEL = load_model(ckpt_path, device=str(_DEVICE), num_classes=len(classes))
    return _MODEL, _DEVICE


def get_expression(img):
    """
    img: can be numpy array (RGB), PIL.Image, or filepath string
    returns: expression string
    """
    model, device = _load_cached_model()

    img_bgr = _to_bgr_np(img)
    if img_bgr is None:
        return "UNKNOWN"

    face = detected_face(img_bgr)
    if face == 0:
        # fallback: use full frame grayscale
        face = Image.fromarray(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY))

    x = trnscm(face).unsqueeze(0).to(device)  # [1,3,H,W] if trnscm ensures RGB
    with torch.no_grad():
        logits = model(x)
        pred = int(torch.argmax(logits, dim=1).item())

    return classes.get(pred, "UNKNOWN").capitalize()