import os
import urllib.request
import numpy as np
import cv2
from pathlib import Path
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input

BASE_DIR = Path(__file__).resolve().parent.parent
FACE_DETECTOR_DIR = BASE_DIR / "models" / "face_detector"

FACE_DETECTOR_FILES = {
    "deploy.prototxt": (
        "https://raw.githubusercontent.com/opencv/opencv/master/"
        "samples/dnn/face_detector/deploy.prototxt"
    ),
    "res10_300x300_ssd_iter_140000.caffemodel": (
        "https://raw.githubusercontent.com/opencv/opencv_3rdparty/"
        "dnn_samples_face_detector_20170830/"
        "res10_300x300_ssd_iter_140000.caffemodel"
    ),
}

CLASSES = ["with_mask", "without_mask"]


def download_face_detector(dest_dir: Path = FACE_DETECTOR_DIR) -> None:
    dest_dir.mkdir(parents=True, exist_ok=True)
    for filename, url in FACE_DETECTOR_FILES.items():
        dest = dest_dir / filename
        if dest.exists():
            print(f"[OK] {filename} already exists, skipping download.")
            continue
        print(f"[DOWNLOADING] {filename} ...")
        urllib.request.urlretrieve(url, dest)
        print(f"[DONE] Saved to {dest}")


def load_face_detector(model_dir: Path = FACE_DETECTOR_DIR):
    prototxt = str(model_dir / "deploy.prototxt")
    caffemodel = str(model_dir / "res10_300x300_ssd_iter_140000.caffemodel")
    if not os.path.exists(prototxt) or not os.path.exists(caffemodel):
        raise FileNotFoundError(
            "Face detector model files not found. "
            "Run: python src/utils.py --download-face-detector"
        )
    return cv2.dnn.readNet(prototxt, caffemodel)


def detect_faces(frame: np.ndarray, net, confidence_threshold: float = 0.5):
    """
    Returns list of (startX, startY, endX, endY) for each face found.
    Adds 5% padding around each detected face region.
    """
    h, w = frame.shape[:2]
    blob = cv2.dnn.blobFromImage(
        cv2.resize(frame, (300, 300)), 1.0, (300, 300),
        (104.0, 177.0, 123.0)
    )
    net.setInput(blob)
    detections = net.forward()

    faces = []
    for i in range(detections.shape[2]):
        confidence = detections[0, 0, i, 2]
        if confidence < confidence_threshold:
            continue
        box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
        startX, startY, endX, endY = box.astype(int)

        # Add 5% padding for better crop context
        pad_x = int((endX - startX) * 0.05)
        pad_y = int((endY - startY) * 0.05)
        startX = max(0, startX - pad_x)
        startY = max(0, startY - pad_y)
        endX = min(w, endX + pad_x)
        endY = min(h, endY + pad_y)

        faces.append((startX, startY, endX, endY))
    return faces


def preprocess_face(face_roi: np.ndarray, target_size: tuple = (224, 224)) -> np.ndarray:
    face = cv2.resize(face_roi, target_size)
    face = cv2.cvtColor(face, cv2.COLOR_BGR2RGB)
    face = preprocess_input(face.astype("float32"))   # → [-1, 1] for MobileNetV2
    return np.expand_dims(face, axis=0)


def get_label_color(label: str) -> tuple:
    return (0, 255, 0) if label == "with_mask" else (0, 0, 255)


def load_dataset_paths(dataset_dir: Path):
    """
    Returns (image_paths, labels) lists from a folder structured as:
        dataset_dir/
            with_mask/
            without_mask/
    """
    image_paths, labels = [], []
    for class_name in CLASSES:
        class_dir = dataset_dir / class_name
        if not class_dir.exists():
            print(f"[WARNING] Directory not found: {class_dir}")
            continue
        for ext in ("*.jpg", "*.jpeg", "*.png"):
            for img_path in class_dir.glob(ext):
                image_paths.append(str(img_path))
                labels.append(class_name)
    return image_paths, labels


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--download-face-detector", action="store_true")
    args = parser.parse_args()

    if args.download_face_detector:
        download_face_detector()