import base64
import threading
import time
import uuid
from collections import deque
from pathlib import Path

import cv2
import numpy as np
import mediapipe as mp
from mediapipe.tasks import python as mp_python
from mediapipe.tasks.python import vision as mp_vision


BASE_DIR = Path(__file__).resolve().parent
MODEL_PATH = BASE_DIR / "face_landmarker.task"
OBJECT_MODEL_PATH = BASE_DIR / "efficientdet_lite0.tflite"

# Detection thresholds (tune against recorded test videos).
PHONE_SCORE = 0.40
PERSON_SCORE = 0.50
OBJECT_MODEL_SCORE = 0.30

# Temporal verification: a detection must be present in this many
# recent analyzed frames before a verified event is raised.
HISTORY_SIZE = 10
MIN_PHONE_HITS = 6
MIN_EXTRA_PERSON_HITS = 6
EVENT_COOLDOWN_SECONDS = 5.0

# Approximate face direction thresholds, used only as a fallback
# warning indicator (not as proof of gaze or misconduct).
HEAD_YAW_LIMIT_DEG = 28.0
HEAD_PITCH_LIMIT_DEG = 25.0

# Reject obviously malformed/empty images without raising to callers.
MAX_IMAGE_BYTES = 8 * 1024 * 1024

for model_path, help_text in (
    (MODEL_PATH, "Download face_landmarker.task and put it next to vision.py."),
    (OBJECT_MODEL_PATH, "Download efficientdet_lite0.tflite and put it next to vision.py."),
):
    if not model_path.is_file():
        raise FileNotFoundError(
            f"Model file not found (or it is a folder): {model_path}\n{help_text}"
        )

_face_options = mp_vision.FaceLandmarkerOptions(
    base_options=mp_python.BaseOptions(model_asset_path=str(MODEL_PATH)),
    running_mode=mp_vision.RunningMode.IMAGE,
    num_faces=2,
)
_landmarker = mp_vision.FaceLandmarker.create_from_options(_face_options)

_object_options = mp_vision.ObjectDetectorOptions(
    base_options=mp_python.BaseOptions(model_asset_path=str(OBJECT_MODEL_PATH)),
    running_mode=mp_vision.RunningMode.IMAGE,
    score_threshold=OBJECT_MODEL_SCORE,
    category_allowlist=["cell phone", "person"],
    max_results=10,
)
_detector = mp_vision.ObjectDetector.create_from_options(_object_options)

# MediaPipe Tasks objects are shared between requests. Serialize inference.
_model_lock = threading.Lock()

# Session state is protected separately from model inference.
_sessions_lock = threading.Lock()
_sessions = {}


class _TemporalVerifier:
    """Verify a detection across a rolling window of analyzed frames."""

    def __init__(self, window_size=HISTORY_SIZE, min_hits=MIN_PHONE_HITS):
        self.history = deque(maxlen=window_size)
        self.min_hits = min_hits
        self.window_size = window_size

    def update(self, detected):
        self.history.append(bool(detected))
        return (
            len(self.history) == self.window_size
            and sum(self.history) >= self.min_hits
        )

    def reset(self):
        self.history.clear()


class _SessionState:
    def __init__(self):
        self.phone = _TemporalVerifier(HISTORY_SIZE, MIN_PHONE_HITS)
        self.extra_person = _TemporalVerifier(
            HISTORY_SIZE, MIN_EXTRA_PERSON_HITS
        )
        self.last_event_at = {}
        self.last_seen_at = time.monotonic()

    def cooldown_passed(self, event_name, now):
        previous = self.last_event_at.get(event_name, 0.0)
        if now - previous >= EVENT_COOLDOWN_SECONDS:
            self.last_event_at[event_name] = now
            return True
        return False


def _get_session(session_id):
    """Return the state for a session. Caller must hold _sessions_lock."""
    state = _sessions.get(session_id)
    if state is None:
        state = _SessionState()
        _sessions[session_id] = state
    state.last_seen_at = time.monotonic()
    return state


def reset_session(session_id):
    """Call when an interview begins/ends to clear its temporal state."""
    if not session_id:
        return
    with _sessions_lock:
        _sessions.pop(str(session_id), None)


def cleanup_inactive_sessions(max_idle_seconds=3600):
    """Optional housekeeping hook; call periodically from the application."""
    now = time.monotonic()
    with _sessions_lock:
        stale = [
            sid for sid, state in _sessions.items()
            if now - state.last_seen_at > max_idle_seconds
        ]
        for sid in stale:
            _sessions.pop(sid, None)
    return len(stale)


def _decode_data_url(data_url):
    """Decode a base64 data URL or raw base64 string into a MediaPipe image."""
    if not isinstance(data_url, str) or not data_url.strip():
        return None

    encoded = data_url.strip()
    if encoded.startswith("data:"):
        header, separator, payload = encoded.partition(",")
        if not separator or ";base64" not in header.lower():
            return None
        encoded = payload

    # Guard against excessively large input before allocating decoded bytes.
    if len(encoded) > (MAX_IMAGE_BYTES * 4 // 3 + 16):
        return None

    try:
        img_bytes = base64.b64decode(encoded, validate=True)
    except (ValueError, base64.binascii.Error):
        return None

    if not img_bytes or len(img_bytes) > MAX_IMAGE_BYTES:
        return None

    image_array = cv2.imdecode(
        np.frombuffer(img_bytes, dtype=np.uint8), cv2.IMREAD_COLOR
    )
    if image_array is None or image_array.size == 0:
        return None

    rgb = cv2.cvtColor(image_array, cv2.COLOR_BGR2RGB)
    rgb = np.ascontiguousarray(rgb)
    return mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb)


def _head_pose_warning(face_landmarks, image_width, image_height):
    """
    Estimate head orientation with solvePnP using selected Face Mesh landmarks.
    This is approximate and depends on camera/model geometry; it is not gaze
    detection. Returns (warning, angles_dict).
    """
    if not face_landmarks:
        return "Your face is not visible!", None

    lm = face_landmarks[0]

    # 2D Face Mesh landmark indices and corresponding generic 3D face points.
    # Approximate generic model points, not personalized calibration.
    landmark_indices = [1, 152, 33, 263, 61, 291]
    model_points = np.array(
        [
            (0.0, 0.0, 0.0),        # nose tip
            (0.0, -63.6, -12.5),    # chin
            (-43.3, 32.7, -26.0),   # left eye outer corner
            (43.3, 32.7, -26.0),    # right eye outer corner
            (-28.9, -28.9, -24.1),  # left mouth corner
            (28.9, -28.9, -24.1),   # right mouth corner
        ],
        dtype=np.float64,
    )
    image_points = np.array(
        [
            (lm[index].x * image_width, lm[index].y * image_height)
            for index in landmark_indices
        ],
        dtype=np.float64,
    )

    focal_length = float(image_width)
    center = (image_width / 2.0, image_height / 2.0)
    camera_matrix = np.array(
        [
            [focal_length, 0.0, center[0]],
            [0.0, focal_length, center[1]],
            [0.0, 0.0, 1.0],
        ],
        dtype=np.float64,
    )
    distortion = np.zeros((4, 1), dtype=np.float64)

    try:
        success, rotation_vector, _ = cv2.solvePnP(
            model_points,
            image_points,
            camera_matrix,
            distortion,
            flags=cv2.SOLVEPNP_ITERATIVE,
        )
        if not success:
            return "", None

        rotation_matrix, _ = cv2.Rodrigues(rotation_vector)
        angles = cv2.RQDecomp3x3(rotation_matrix)[0]
        pitch, yaw, roll = (float(angles[0]), float(angles[1]), float(angles[2]))
        angle_info = {"pitch": pitch, "yaw": yaw, "roll": roll}

        if abs(yaw) > HEAD_YAW_LIMIT_DEG:
            return "Your head is turned away from the camera. Please face the screen.", angle_info
        if abs(pitch) > HEAD_PITCH_LIMIT_DEG:
            return "Your head is tilted up or down. Please face the screen.", angle_info

        return "", angle_info
    except (cv2.error, ValueError, IndexError, TypeError):
        # Do not fail the entire interview if pose estimation fails.
        return "", None


def _legacy_face_warning(faces, image_width=None, image_height=None):
    """Fallback warning based on the original normalized landmark ratios."""
    if not faces:
        return "Your face is not visible!"
    if len(faces) > 1:
        return "More than one face is visible in the frame!"

    lm = faces[0]
    try:
        nose, left, right, top, chin = lm[1], lm[234], lm[454], lm[10], lm[152]
        h = (nose.x - left.x) / (right.x - left.x + 1e-6)
        v = (nose.y - top.y) / (chin.y - top.y + 1e-6)

        if h < 0.35 or h > 0.65:
            return "Your face is turned away. Please face the camera!"
        if v < 0.40 or v > 0.75:
            return "Please keep your face directed toward the screen."
    except (IndexError, ZeroDivisionError, TypeError):
        return ""
    return ""


def inspect(data_url: str, session_id="default"):
    """
    Analyze one frame.

    Returns (violation, warning), preserving the original interface.
    violation is one of: "phone", "multiple_people", or "".
    A violation is only returned after persistent detections across a window.
    `session_id` must uniquely identify the active interview session.
    """
    mp_image = _decode_data_url(data_url)
    if mp_image is None:
        return "", ""

    width, height = mp_image.width, mp_image.height

    try:
        with _model_lock:
            face_result = _landmarker.detect(mp_image)
            object_result = _detector.detect(mp_image)
    except (RuntimeError, ValueError):
        # Avoid crashing a request because an inference frame was malformed.
        return "", ""

    faces = face_result.face_landmarks or []
    phone_seen = False
    persons = 0
    phone_scores = []
    person_scores = []

    for detection in (object_result.detections or []):
        for category in (detection.categories or []):
            name = (category.category_name or "").lower()
            score = float(category.score or 0.0)
            if name == "cell phone" and score >= PHONE_SCORE:
                phone_seen = True
                phone_scores.append(score)
            elif name == "person" and score >= PERSON_SCORE:
                persons += 1
                person_scores.append(score)

    # More than one detected face is strong extra-face evidence.
    # Person detections are a separate signal; they may overlap with faces.
    extra_person_seen = len(faces) > 1 or persons > 1

    now = time.monotonic()
    with _sessions_lock:
        state = _get_session(str(session_id or "default"))
        verified_phone = state.phone.update(phone_seen)
        verified_extra_person = state.extra_person.update(extra_person_seen)

        violation = ""
        if verified_phone and state.cooldown_passed("phone", now):
            violation = "phone"
        elif verified_extra_person and state.cooldown_passed("multiple_people", now):
            violation = "multiple_people"

    if not faces:
        warning = "Your face is not visible!"
        pose_angles = None
    elif len(faces) > 1:
        warning = "More than one face is visible in the frame!"
        pose_angles = None
    else:
        warning, pose_angles = _head_pose_warning(faces, width, height)
        if not warning:
            warning = _legacy_face_warning(faces, width, height)

    return violation, warning


def analyze(data_url: str, session_id="default") -> str:
    """Return a warning string, keeping compatibility with existing callers."""
    return inspect(data_url, session_id=session_id)[1]


def detect_violation(data_url: str, session_id="default") -> str:
    """Return 'phone', 'multiple_people', or '' after temporal verification."""
    return inspect(data_url, session_id=session_id)[0]


def inspect_detailed(data_url: str, session_id="default"):
    """
    Detailed alternative API for new integrations.
    Returns a JSON-serializable dict; no image data is retained.
    """
    mp_image = _decode_data_url(data_url)
    if mp_image is None:
        return {
            "ok": False,
            "session_id": str(session_id or "default"),
            "timestamp": time.time(),
            "violation": "",
            "warning": "",
            "reason": "invalid_or_unsupported_image",
        }

    # Re-encode to data URL to reuse the single inference implementation.
    # This avoids maintaining two separate implementations of detection logic.
    # JPEG is adequate for analysis frames; callers should use inspect() directly
    # when they only need the legacy return shape.
    try:
        rgb = mp_image.numpy_view()
        bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
        success, encoded = cv2.imencode(".jpg", bgr)
        if not success:
            return {
                "ok": False,
                "session_id": str(session_id or "default"),
                "timestamp": time.time(),
                "violation": "",
                "warning": "",
                "reason": "image_encode_failed",
            }
        data_url = "data:image/jpeg;base64," + base64.b64encode(encoded).decode("ascii")
    except (cv2.error, ValueError):
        return {
            "ok": False,
            "session_id": str(session_id or "default"),
            "timestamp": time.time(),
            "violation": "",
            "warning": "",
            "reason": "image_processing_failed",
        }

    violation, warning = inspect(data_url, session_id=session_id)
    return {
        "ok": True,
        "session_id": str(session_id or "default"),
        "timestamp": time.time(),
        "violation": violation,
        "warning": warning,
        "event_id": str(uuid.uuid4()) if violation else None,
    }