import cv2
import numpy as np
import os
import urllib.request
from huggingface_hub import hf_hub_download

REPO_ID = "codernotme/kataria_optical"
MODEL_FILE = "lbfmodel.yaml"

_facemark = None
_face_cascade = None

def ensure_model():
    global MODEL_FILE
    if not os.path.exists(MODEL_FILE):
        print(f"Attempting to download {MODEL_FILE} from HF Hub...")
        try:
            downloaded_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILE)
            MODEL_FILE = downloaded_path
            print("Download from HF Hub complete.")
        except Exception as e:
            print(f"HF Hub download failed ({e}), falling back to GitHub...")
            MODEL_URL = "https://raw.githubusercontent.com/kurnianggoro/GSOC2017/master/data/lbfmodel.yaml"
            try:
                urllib.request.urlretrieve(MODEL_URL, MODEL_FILE)
                print("Download from GitHub complete.")
            except Exception as e2:
                raise RuntimeError(f"Failed to download LBF model from all sources: {e2}")


def _get_facemark():
    global _facemark
    if _facemark is None:
        ensure_model()
        _facemark = cv2.face.createFacemarkLBF()
        try:
            _facemark.loadModel(MODEL_FILE)
        except Exception as e:
            raise RuntimeError(f"Could not load LBF model: {e}")
    return _facemark


def _get_face_cascade():
    global _face_cascade
    if _face_cascade is None:
        face_cascade_path = cv2.data.haarcascades + "haarcascade_frontalface_alt2.xml"
        _face_cascade = cv2.CascadeClassifier(face_cascade_path)
        if _face_cascade.empty():
            _face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_frontalface_default.xml")
    return _face_cascade


def get_landmarks(image_input):
    """
    Extracts 68 face landmarks using OpenCV FacemarkLBF.
    Returns a list/array of (x, y) coordinates.
    """
    facemark = _get_facemark()

    if isinstance(image_input, str):
        image = cv2.imread(image_input)
        if image is None:
            raise ValueError(f"Could not read image: {image_input}")
    else:
        if hasattr(image_input, "mode"):
            image = cv2.cvtColor(np.array(image_input), cv2.COLOR_RGB2BGR)
        else:
            image = np.array(image_input)
            if image.ndim == 3 and image.shape[2] == 3:
                image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        if image is None or image.size == 0:
            raise ValueError("Could not read image data.")
    
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    
    # Detect faces using Haar Cascade (included in cv2)
    face_cascade = _get_face_cascade()
    faces = face_cascade.detectMultiScale(gray, 1.3, 5)
    
    if len(faces) == 0:
        raise ValueError("No face detected.")

    # Find landmarks
    ok, landmarks = facemark.fit(gray, faces)
    if not ok:
        raise ValueError("Could not detect landmarks.")
        
    # landmarks comes as a list of numpy arrays, we extract the first face's landmarks
    # landmarks[0] is shape (1, 68, 2)
    return landmarks[0][0]