"""Pure skeleton drawing functions using OpenCV.

All functions draw directly onto a BGR frame (numpy array) in-place.
Colors are in BGR format for OpenCV.
"""

from __future__ import annotations

import numpy as np

_INT32_MAX = 2**31 - 1
_INT32_MIN = -(2**31)

# ============================================================================
# Connection Constants
# ============================================================================

# MediaPipe Pose 33-landmark body connections
MEDIAPIPE_BODY_CONNECTIONS: list[tuple[int, int]] = [
    (0, 1),
    (1, 2),
    (2, 3),
    (3, 7),  # Right eye
    (0, 4),
    (4, 5),
    (5, 6),
    (6, 8),  # Left eye
    (9, 10),  # Mouth
    (11, 12),  # Shoulders
    (11, 13),
    (13, 15),  # Left arm
    (12, 14),
    (14, 16),  # Right arm
    (15, 17),
    (15, 19),
    (15, 21),
    (17, 19),  # Left wrist
    (16, 18),
    (16, 20),
    (16, 22),
    (18, 20),  # Right wrist
    (11, 23),
    (12, 24),
    (23, 24),  # Torso
    (23, 25),
    (25, 27),
    (27, 29),
    (27, 31),
    (29, 31),  # Left leg
    (24, 26),
    (26, 28),
    (28, 30),
    (28, 32),
    (30, 32),  # Right leg
]

# 21-joint hand connections (shared by MediaPipe hands and WiLoR)
HAND_CONNECTIONS: list[tuple[int, int]] = [
    (0, 1),
    (1, 2),
    (2, 3),
    (3, 4),  # Thumb
    (0, 5),
    (5, 6),
    (6, 7),
    (7, 8),  # Index
    (0, 9),
    (9, 10),
    (10, 11),
    (11, 12),  # Middle
    (0, 13),
    (13, 14),
    (14, 15),
    (15, 16),  # Ring
    (0, 17),
    (17, 18),
    (18, 19),
    (19, 20),  # Pinky
    (5, 9),
    (9, 13),
    (13, 17),  # Palm
]

# COCO 17-body-joint connections
# Joint order: 0=nose, 1=L_eye, 2=R_eye, 3=L_ear, 4=R_ear,
#   5=L_shoulder, 6=R_shoulder, 7=L_elbow, 8=R_elbow,
#   9=L_wrist, 10=R_wrist, 11=L_hip, 12=R_hip,
#   13=L_knee, 14=R_knee, 15=L_ankle, 16=R_ankle
# COCO upper-body connections only (no legs 13-16, no feet 17-22)
COCO_UPPER_BODY_CONNECTIONS: list[tuple[int, int]] = [
    (0, 1),
    (0, 2),
    (1, 3),
    (2, 4),  # Head
    (5, 6),  # Shoulders
    (5, 7),
    (7, 9),  # Left arm
    (6, 8),
    (8, 10),  # Right arm
    (5, 11),
    (6, 12),
    (11, 12),  # Torso
]

# Full COCO 17-body connections (upper body + legs)
COCO_BODY_CONNECTIONS: list[tuple[int, int]] = COCO_UPPER_BODY_CONNECTIONS + [
    (11, 13),
    (13, 15),  # Left leg
    (12, 14),
    (14, 16),  # Right leg
]

# SMPL-X 22-body-joint connections
# Joint order: 0=pelvis, 1=L_hip, 2=R_hip, 3=spine1, 4=L_knee, 5=R_knee,
#   6=spine2, 7=L_ankle, 8=R_ankle, 9=spine3, 10=L_foot, 11=R_foot,
#   12=neck, 13=L_collar, 14=R_collar, 15=head,
#   16=L_shoulder, 17=R_shoulder, 18=L_elbow, 19=R_elbow,
#   20=L_wrist, 21=R_wrist
SMPLX_BODY_CONNECTIONS: list[tuple[int, int]] = [
    # Spine: pelvis -> spine1 -> spine2 -> spine3 -> neck -> head
    (0, 3),
    (3, 6),
    (6, 9),
    (9, 12),
    (12, 15),
    # Hip bar: pelvis -> hips
    (0, 1),
    (0, 2),
    # Left arm: spine3 -> L_collar -> L_shoulder -> L_elbow -> L_wrist
    (9, 13),
    (13, 16),
    (16, 18),
    (18, 20),
    # Right arm: spine3 -> R_collar -> R_shoulder -> R_elbow -> R_wrist
    (9, 14),
    (14, 17),
    (17, 19),
    (19, 21),
    # Left leg: L_hip -> L_knee -> L_ankle -> L_foot
    (1, 4),
    (4, 7),
    (7, 10),
    # Right leg: R_hip -> R_knee -> R_ankle -> R_foot
    (2, 5),
    (5, 8),
    (8, 11),
]

# SMPL-X hand joints: 15 per hand, NO wrist (wrist is body joint 20/21).
# Order: index(0-2), middle(3-5), pinky(6-8), ring(9-11), thumb(12-14)
# Each triplet is MCP -> PIP -> DIP.  Intra-finger chains only;
# wrist-to-MCP connections are drawn separately using body wrist joint.
SMPLX_HAND_CONNECTIONS: list[tuple[int, int]] = [
    (0, 1),
    (1, 2),  # Index: MCP -> PIP -> DIP
    (3, 4),
    (4, 5),  # Middle: MCP -> PIP -> DIP
    (6, 7),
    (7, 8),  # Pinky: MCP -> PIP -> DIP
    (9, 10),
    (10, 11),  # Ring: MCP -> PIP -> DIP
    (12, 13),
    (13, 14),  # Thumb: CMC -> IP -> TIP
]

# MCP (base) joint local indices for connecting to wrist
SMPLX_HAND_MCP_INDICES: list[int] = [0, 3, 6, 9, 12]  # index, middle, pinky, ring, thumb

# ============================================================================
# Color Scheme (BGR for OpenCV)
# ============================================================================

COLOR_BODY = (250, 165, 96)  # Blue-ish (matches PoseViewer2D #60a5fa)
COLOR_LEFT_HAND = (74, 222, 128)  # Green (matches #4ade80)
COLOR_RIGHT_HAND = (168, 85, 247)  # Purple (matches #f755a8 -> BGR)
COLOR_JOINT = (255, 255, 255)  # White joint dots
COLOR_FACE = (0, 200, 255)  # Orange/amber for face box
COLOR_FACE_LANDMARK = (0, 220, 255)  # Slightly brighter amber for landmarks


def _line_thickness(h: int, w: int) -> int:
    """Scale line thickness with frame resolution."""
    return max(1, int(min(h, w) / 300))


def _circle_radius(h: int, w: int) -> int:
    """Scale joint circle radius with frame resolution."""
    return max(1, int(min(h, w) / 200))


def _draw_connections(
    frame: np.ndarray,
    keypoints: np.ndarray,
    connections: list[tuple[int, int]],
    color: tuple[int, int, int],
    thickness: int,
    confidence: np.ndarray | None = None,
    min_confidence: float = 0.1,
) -> None:
    """Draw skeleton connections on a frame."""
    import cv2

    n_kpts = len(keypoints)
    for i, j in connections:
        if i >= n_kpts or j >= n_kpts:
            continue

        # Skip non-finite coordinates (NaN, Inf)
        if not (np.isfinite(keypoints[i][0]) and np.isfinite(keypoints[i][1])):
            continue
        if not (np.isfinite(keypoints[j][0]) and np.isfinite(keypoints[j][1])):
            continue

        # Clamp to int32 range for OpenCV compatibility
        x1 = max(_INT32_MIN, min(_INT32_MAX, int(keypoints[i][0])))
        y1 = max(_INT32_MIN, min(_INT32_MAX, int(keypoints[i][1])))
        x2 = max(_INT32_MIN, min(_INT32_MAX, int(keypoints[j][0])))
        y2 = max(_INT32_MIN, min(_INT32_MAX, int(keypoints[j][1])))

        # Skip zero-coordinate points (no detection)
        if (x1 == 0 and y1 == 0) or (x2 == 0 and y2 == 0):
            continue

        # Skip low-confidence joints
        if confidence is not None:
            if i < len(confidence) and confidence[i] < min_confidence:
                continue
            if j < len(confidence) and confidence[j] < min_confidence:
                continue

        cv2.line(frame, (x1, y1), (x2, y2), color, thickness, cv2.LINE_AA)


def _draw_joints(
    frame: np.ndarray,
    keypoints: np.ndarray,
    radius: int,
    color: tuple[int, int, int] = COLOR_JOINT,
    confidence: np.ndarray | None = None,
    min_confidence: float = 0.1,
) -> None:
    """Draw joint circles on a frame."""
    import cv2

    for idx, kpt in enumerate(keypoints):
        if not (np.isfinite(kpt[0]) and np.isfinite(kpt[1])):
            continue
        x = max(_INT32_MIN, min(_INT32_MAX, int(kpt[0])))
        y = max(_INT32_MIN, min(_INT32_MAX, int(kpt[1])))
        if x == 0 and y == 0:
            continue
        if confidence is not None and idx < len(confidence) and confidence[idx] < min_confidence:
            continue
        cv2.circle(frame, (x, y), radius, color, -1, cv2.LINE_AA)


# ============================================================================
# Public Drawing Functions
# ============================================================================


def draw_mediapipe_skeleton(
    frame: np.ndarray,
    keypoints: np.ndarray,
    confidence: np.ndarray | None = None,
    hands_l: np.ndarray | None = None,
    hands_r: np.ndarray | None = None,
) -> None:
    """Draw MediaPipe skeleton on a video frame.

    Args:
        frame: BGR video frame (H, W, 3), modified in-place.
        keypoints: Body landmarks shape (33, 2|3). If values < 2, treated as
            normalized 0-1 and scaled to frame dimensions.
        confidence: Per-joint confidence (33,). Joints below 0.1 are hidden.
        hands_l: Left hand landmarks (21, 2|3). Same normalization rules.
        hands_r: Right hand landmarks (21, 2|3). Same normalization rules.
    """
    h, w = frame.shape[:2]
    thickness = _line_thickness(h, w)
    radius = _circle_radius(h, w)

    def to_pixels(kpts: np.ndarray) -> np.ndarray:
        """Convert normalized coords to pixel coords if needed."""
        kpts = np.asarray(kpts, dtype=np.float64)
        result = kpts[:, :2].copy()
        # Use nanmax to handle NaN values; check if all finite values are in [0, 1]
        finite_mask = np.isfinite(result)
        if finite_mask.any():
            max_val = np.nanmax(np.abs(result[finite_mask]))
            if max_val <= 1.0:
                result[:, 0] *= w
                result[:, 1] *= h
        return result

    body_px = to_pixels(keypoints)

    # Draw body
    _draw_connections(frame, body_px, MEDIAPIPE_BODY_CONNECTIONS, COLOR_BODY, thickness, confidence)
    _draw_joints(frame, body_px, radius, COLOR_BODY, confidence)

    # Draw left hand
    if hands_l is not None and len(hands_l) > 0:
        lh_px = to_pixels(hands_l)
        _draw_connections(frame, lh_px, HAND_CONNECTIONS, COLOR_LEFT_HAND, thickness)
        _draw_joints(frame, lh_px, radius, COLOR_LEFT_HAND)

    # Draw right hand
    if hands_r is not None and len(hands_r) > 0:
        rh_px = to_pixels(hands_r)
        _draw_connections(frame, rh_px, HAND_CONNECTIONS, COLOR_RIGHT_HAND, thickness)
        _draw_joints(frame, rh_px, radius, COLOR_RIGHT_HAND)


def draw_wilor_hands(
    frame: np.ndarray,
    kpts_2d: np.ndarray,
    right_flags: np.ndarray | None = None,
) -> None:
    """Draw WiLoR hand detections on a video frame.

    Args:
        frame: BGR video frame (H, W, 3), modified in-place.
        kpts_2d: 2D keypoints shape (N, 21, 2) in pixel coords.
        right_flags: Boolean array (N,) -- True means right hand.
            If None, alternates left/right.
    """
    h, w = frame.shape[:2]
    thickness = _line_thickness(h, w)
    radius = _circle_radius(h, w)

    for det_idx in range(len(kpts_2d)):
        hand_kpts = kpts_2d[det_idx]  # (21, 2)
        is_right = right_flags[det_idx] if right_flags is not None else (det_idx % 2 == 1)
        color = COLOR_RIGHT_HAND if is_right else COLOR_LEFT_HAND

        _draw_connections(frame, hand_kpts, HAND_CONNECTIONS, color, thickness)
        _draw_joints(frame, hand_kpts, radius, color)


def draw_nlf_body(
    frame: np.ndarray,
    joints2d: np.ndarray,
) -> None:
    """Draw NLF/SMPL-X body skeleton on a video frame.

    Args:
        frame: BGR video frame (H, W, 3), modified in-place.
        joints2d: 2D joint positions shape (N, 55, 2) in pixel coords.
            For each detection: 22 body + 15 left hand + 15 right hand + 3 extra.
    """
    h, w = frame.shape[:2]
    thickness = _line_thickness(h, w)
    radius = _circle_radius(h, w)

    for det_idx in range(len(joints2d)):
        joints = joints2d[det_idx]  # (55, 2)

        # SMPL-X 55-joint layout:
        #   0-21:  22 body joints (20=L_wrist, 21=R_wrist)
        #   22:    jaw
        #   23:    left eye
        #   24:    right eye
        #   25-39: 15 left hand joints  (index, middle, pinky, ring, thumb × 3)
        #   40-54: 15 right hand joints (same order)

        # Body joints (0-21)
        body_joints = joints[:22]
        _draw_connections(frame, body_joints, SMPLX_BODY_CONNECTIONS, COLOR_BODY, thickness)
        _draw_joints(frame, body_joints, radius, COLOR_BODY)

        # Jaw / eyes (22-24) -- draw as dots only
        extra_joints = joints[22:25]
        if len(extra_joints) > 0:
            _draw_joints(frame, extra_joints, radius, COLOR_BODY)

        # Left hand (25-39, mapped to 0-14 for SMPLX_HAND_CONNECTIONS)
        left_hand = joints[25:40]
        if len(left_hand) == 15:
            _draw_connections(frame, left_hand, SMPLX_HAND_CONNECTIONS, COLOR_LEFT_HAND, thickness)
            _draw_joints(frame, left_hand, radius, COLOR_LEFT_HAND)
            # Connect wrist (body joint 20) to each finger's MCP
            l_wrist = joints[20:21]  # (1, 2)
            for mcp_idx in SMPLX_HAND_MCP_INDICES:
                mcp = left_hand[mcp_idx : mcp_idx + 1]  # (1, 2)
                _draw_connections(
                    frame,
                    np.concatenate([l_wrist, mcp], axis=0),
                    [(0, 1)],
                    COLOR_LEFT_HAND,
                    thickness,
                )

        # Right hand (40-54, mapped to 0-14)
        right_hand = joints[40:55]
        if len(right_hand) == 15:
            _draw_connections(
                frame,
                right_hand,
                SMPLX_HAND_CONNECTIONS,
                COLOR_RIGHT_HAND,
                thickness,
            )
            _draw_joints(frame, right_hand, radius, COLOR_RIGHT_HAND)
            # Connect wrist (body joint 21) to each finger's MCP
            r_wrist = joints[21:22]  # (1, 2)
            for mcp_idx in SMPLX_HAND_MCP_INDICES:
                mcp = right_hand[mcp_idx : mcp_idx + 1]  # (1, 2)
                _draw_connections(
                    frame,
                    np.concatenate([r_wrist, mcp], axis=0),
                    [(0, 1)],
                    COLOR_RIGHT_HAND,
                    thickness,
                )


def draw_teaser_face(
    frame: np.ndarray,
    boxes: np.ndarray,
    scores: np.ndarray | None = None,
    jaw_pose: np.ndarray | None = None,
    expression: np.ndarray | None = None,
    eyelid_params: np.ndarray | None = None,
) -> None:
    """Draw TEASER face detection boxes and FLAME parameter HUD.

    Args:
        frame: BGR video frame (H, W, 3), modified in-place.
        boxes: Bounding boxes (N, 4) in xyxy pixel coords.
        scores: Detection confidence (N,).
        jaw_pose: FLAME jaw pose (N, 3).
        expression: FLAME expression coefficients (N, 50).
        eyelid_params: FLAME eyelid parameters (N, 2).
    """
    import cv2

    h, w = frame.shape[:2]
    thickness = _line_thickness(h, w)
    font_scale = max(0.3, min(h, w) / 800)
    font = cv2.FONT_HERSHEY_SIMPLEX

    for i in range(len(boxes)):
        box = boxes[i]
        if not np.isfinite(box).all():
            continue

        x1, y1, x2, y2 = int(box[0]), int(box[1]), int(box[2]), int(box[3])

        # Draw bounding box
        cv2.rectangle(frame, (x1, y1), (x2, y2), COLOR_FACE, thickness)

        # Build text lines for HUD
        lines: list[str] = []
        if scores is not None and i < len(scores):
            lines.append(f"conf: {float(scores[i]):.2f}")
        if jaw_pose is not None and i < len(jaw_pose):
            jp = jaw_pose[i]
            lines.append(f"jaw: {jp[0]:.2f} {jp[1]:.2f} {jp[2]:.2f}")
        if eyelid_params is not None and i < len(eyelid_params):
            ep = eyelid_params[i]
            lines.append(f"eyelid: {ep[0]:.2f} {ep[1]:.2f}")
        if expression is not None and i < len(expression):
            # Show L2 norm of expression vector as a compact summary
            expr_norm = float(np.linalg.norm(expression[i]))
            lines.append(f"expr|: {expr_norm:.1f}")

        # Draw text above the box
        text_y = max(y1 - 4, 12)
        for line in reversed(lines):
            (tw, th), _ = cv2.getTextSize(line, font, font_scale, 1)
            # Background rectangle for readability
            cv2.rectangle(
                frame,
                (x1, text_y - th - 4),
                (x1 + tw + 4, text_y + 2),
                (0, 0, 0),
                -1,
            )
            cv2.putText(frame, line, (x1 + 2, text_y), font, font_scale, COLOR_FACE, 1, cv2.LINE_AA)
            text_y -= th + 6


# 68-point FAN landmark connections (face contour, brows, nose, eyes, mouth)
FLAME_FACE_CONNECTIONS: list[tuple[int, int]] = [
    # Jaw contour (0-16)
    *[(i, i + 1) for i in range(16)],
    # Right eyebrow (17-21)
    *[(i, i + 1) for i in range(17, 21)],
    # Left eyebrow (22-26)
    *[(i, i + 1) for i in range(22, 26)],
    # Nose bridge (27-30)
    *[(i, i + 1) for i in range(27, 30)],
    # Nose bottom (31-35)
    *[(i, i + 1) for i in range(31, 35)],
    # Right eye (36-41, closed loop)
    *[(i, i + 1) for i in range(36, 41)],
    (41, 36),
    # Left eye (42-47, closed loop)
    *[(i, i + 1) for i in range(42, 47)],
    (47, 42),
    # Outer mouth (48-59, closed loop)
    *[(i, i + 1) for i in range(48, 59)],
    (59, 48),
    # Inner mouth (60-67, closed loop)
    *[(i, i + 1) for i in range(60, 67)],
    (67, 60),
]


def draw_teaser_landmarks(
    frame: np.ndarray,
    landmarks_2d: np.ndarray,
) -> None:
    """Draw 68-point FAN face landmarks from FLAME on a video frame.

    Args:
        frame: BGR video frame (H, W, 3), modified in-place.
        landmarks_2d: 2D landmarks shape (N, 68, 2) in pixel coords.
            N is the number of face detections for this frame.
    """
    h, w = frame.shape[:2]
    thickness = _line_thickness(h, w)
    radius = max(1, _circle_radius(h, w) // 2)  # Smaller dots for face

    for det_idx in range(len(landmarks_2d)):
        lmk = landmarks_2d[det_idx]  # (68, 2)
        _draw_connections(frame, lmk, FLAME_FACE_CONNECTIONS, COLOR_FACE_LANDMARK, thickness)
        _draw_joints(frame, lmk, radius, COLOR_FACE_LANDMARK)


def draw_rtmpose_wholebody(
    frame: np.ndarray,
    keypoints: np.ndarray,
) -> None:
    """Draw COCO-WholeBody skeleton on a video frame.

    Args:
        frame: BGR video frame (H, W, 3), modified in-place.
        keypoints: Whole-body landmarks shape ``(N, 133, 3)`` in pixel coords.
            Layout: body(0-16) + foot(17-22) + face(23-90) + left_hand(91-111)
            + right_hand(112-132).  Channel 2 is confidence.
    """
    h, w = frame.shape[:2]
    thickness = _line_thickness(h, w)
    radius = _circle_radius(h, w)
    face_radius = max(1, radius // 2)

    for det_idx in range(len(keypoints)):
        kpts = keypoints[det_idx]  # (133, 3)
        conf = kpts[:, 2]

        # Upper body only (0-12): head, arms, torso — skip legs (13-16) and feet (17-22)
        body_xy = kpts[:17, :2]
        _draw_connections(
            frame,
            body_xy,
            COCO_UPPER_BODY_CONNECTIONS,
            COLOR_BODY,
            thickness,
            conf[:17],
        )
        _draw_joints(frame, body_xy[:13], radius, COLOR_BODY, conf[:13])

        # Face (23-90, 68 points) — use FLAME connections (same 68-point layout)
        face_xy = kpts[23:91, :2]
        _draw_connections(
            frame,
            face_xy,
            FLAME_FACE_CONNECTIONS,
            COLOR_FACE_LANDMARK,
            thickness,
        )
        _draw_joints(frame, face_xy, face_radius, COLOR_FACE_LANDMARK)

        # Left hand (91-111, 21 points)
        left_hand = kpts[91:112, :2]
        _draw_connections(
            frame,
            left_hand,
            HAND_CONNECTIONS,
            COLOR_LEFT_HAND,
            thickness,
        )
        _draw_joints(frame, left_hand, radius, COLOR_LEFT_HAND)

        # Right hand (112-132, 21 points)
        right_hand = kpts[112:133, :2]
        _draw_connections(
            frame,
            right_hand,
            HAND_CONNECTIONS,
            COLOR_RIGHT_HAND,
            thickness,
        )
        _draw_joints(frame, right_hand, radius, COLOR_RIGHT_HAND)
