from __future__ import annotations

"""
modules/ar_webcam.py - Webcam capture + (optional) ArUco + face tracking primitives.

Design:
- Capture runs on a background thread.
- Latest RGB frame is exposed for the OpenGL renderer to upload via glTexSubImage2D.
- ArUco marker pose is computed (when enabled) and published into modules.ar_scene anchors.

This module is "best effort":
- If OpenCV/camera is missing, it degrades gracefully (no crash).
"""

import threading
import time
from dataclasses import dataclass
from typing import Optional, Tuple


try:
    import cv2  # type: ignore
    import numpy as np  # type: ignore

    _CV_OK = True
except Exception:
    cv2 = None
    np = None
    _CV_OK = False


@dataclass
class FramePacket:
    w: int
    h: int
    rgb: "np.ndarray"  # uint8 HxWx3
    ts: float


_lock = threading.Lock()
_latest: Optional[FramePacket] = None
_running = False
_thread: Optional[threading.Thread] = None


def _publish_aruco_anchor(marker_id: int, rvec, tvec) -> None:
    """
    Convert OpenCV rvec/tvec to a rough 4x4 pose matrix and publish to ar_scene.
    Scale is arbitrary (depends on marker size + calibration). For our hologram effect,
    relative pose stability is the main goal.
    """
    try:
        from modules.ar_scene import apply_patch

        R, _ = cv2.Rodrigues(rvec)
        M = np.eye(4, dtype=np.float32)
        M[:3, :3] = R.astype(np.float32)
        M[:3, 3] = np.array(tvec, dtype=np.float32).reshape(3)
        # Row-major list of 16 floats
        key = f"aruco:{int(marker_id)}"
        apply_patch({"anchors": {key: [float(x) for x in M.reshape(-1).tolist()]}})
    except Exception:
        pass


def start(
    *,
    camera_index: int = 0,
    target_fps: int = 30,
    enable_aruco: bool = True,
    aruco_marker_length_m: float = 0.05,
) -> None:
    """
    Starts webcam capture thread (idempotent).

    Notes:
    - For accurate pose you need camera calibration. We use a conservative default
      that still gives stable "table hologram" behavior in practice.
    """
    global _running, _thread
    if _running:
        return
    if not _CV_OK:
        return

    _running = True

    def _loop() -> None:
        global _latest, _running
        cap = None
        try:
            cap = cv2.VideoCapture(int(camera_index), cv2.CAP_DSHOW)
            if not cap.isOpened():
                _running = False
                return

            # Try to stabilize latency
            try:
                cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
            except Exception:
                pass

            # ArUco init
            if enable_aruco:
                try:
                    aruco = cv2.aruco  # type: ignore[attr-defined]
                    dict_ = aruco.getPredefinedDictionary(aruco.DICT_4X4_50)
                    params = aruco.DetectorParameters()
                    detector = aruco.ArucoDetector(dict_, params)
                except Exception:
                    detector = None
            else:
                detector = None

            # Face detection (helmet HUD anchor) - best-effort
            try:
                face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_frontalface_default.xml")
            except Exception:
                face_cascade = None

            # Very rough intrinsics (fallback). If you add calibration later, replace this.
            cam_mtx = None
            dist = None

            sleep_dt = 1.0 / float(max(10, int(target_fps)))
            while _running:
                ok, frame_bgr = cap.read()
                if not ok or frame_bgr is None:
                    time.sleep(0.05)
                    continue

                # BGR -> RGB
                frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
                h, w = frame_rgb.shape[:2]
                pkt = FramePacket(w=w, h=h, rgb=frame_rgb, ts=time.time())
                with _lock:
                    _latest = pkt

                # ArUco detection (best-effort)
                if detector is not None:
                    try:
                        corners, ids, _rej = detector.detectMarkers(frame_bgr)
                        if ids is not None and len(ids) > 0:
                            # Pose estimation requires intrinsics. We'll synthesize if missing.
                            if cam_mtx is None:
                                fx = 0.9 * w
                                fy = 0.9 * w
                                cx = w * 0.5
                                cy = h * 0.5
                                cam_mtx = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float32)
                                dist = np.zeros((5, 1), dtype=np.float32)
                            aruco = cv2.aruco  # type: ignore[attr-defined]
                            rvecs, tvecs, _ = aruco.estimatePoseSingleMarkers(
                                corners, float(aruco_marker_length_m), cam_mtx, dist
                            )
                            for i in range(min(len(ids), len(rvecs))):
                                mid = int(ids[i][0])
                                _publish_aruco_anchor(mid, rvecs[i], tvecs[i])
                    except Exception:
                        pass

                # Face detection (best-effort): publish a simple "face" anchor as normalized center offset
                try:
                    if face_cascade is not None:
                        gray = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2GRAY)
                        faces = face_cascade.detectMultiScale(gray, scaleFactor=1.2, minNeighbors=5, minSize=(60, 60))
                        if faces is not None and len(faces) > 0:
                            x, y, fw, fh = sorted(faces, key=lambda f: f[2] * f[3], reverse=True)[0]
                            cx = (x + fw * 0.5) / float(w)
                            cy = (y + fh * 0.5) / float(h)
                            # Map to [-1,1] offset where (0,0) is screen center.
                            ox = (cx - 0.5) * 2.0
                            oy = (0.5 - cy) * 2.0
                            try:
                                from modules.ar_scene import apply_patch
                                apply_patch({"ui": {"face_offset": {"x": float(ox), "y": float(oy)}}})
                            except Exception:
                                pass
                except Exception:
                    pass

                time.sleep(sleep_dt)
        finally:
            try:
                if cap is not None:
                    cap.release()
            except Exception:
                pass
            _running = False

    _thread = threading.Thread(target=_loop, daemon=True)
    _thread.start()


def stop() -> None:
    global _running
    _running = False


def latest_frame() -> Optional[FramePacket]:
    with _lock:
        return _latest