"""ControlNet-style visual previews for the Tagger tab.

Turns the YOLO pose keypoints produced by ``src.pose_tagger`` into previews a
video/animation pipeline can consume directly:

- ``skeleton_black``   — the pose skeleton on a black, letterboxed canvas
                         (drop-in OpenPose-style ControlNet input)
- ``skeleton_overlay`` — the skeleton drawn on top of the original image
- ``canny``            — Canny-style edge map computed in pure numpy (no
                         OpenCV / model downloads needed)

Keypoints are the COCO-17 arrays returned by ``PoseEstimator.estimate``
(already mapped back into original-image coordinates); rows are
``(x, y, confidence)``. Skeleton edges are confidence-gated at 0.25.
"""
from __future__ import annotations

import os
import tempfile

import numpy as np
from PIL import Image, ImageDraw

# COCO-17 skeleton edges (OpenPose ordering). Indices reference the keypoint
# order of ``src.pose_tagger``: nose, l/r eye, l/r ear, l/r shoulder, l/r
# elbow, l/r wrist, l/r hip, l/r knee, l/r ankle.
_SKELETON_EDGES = [
    (0, 1), (0, 2), (1, 3), (2, 4),           # face
    (5, 6), (5, 7), (7, 9), (6, 8), (8, 10),  # torso + arms
    (5, 11), (6, 12), (11, 12),               # hips
    (11, 13), (13, 15), (12, 14), (14, 16),   # legs
]

_KP_CONF = 0.25

_PALETTE = [
    (255, 64, 128),
    (0, 200, 255),
    (140, 255, 64),
    (255, 170, 0),
    (170, 64, 255),
]

_KINDS = ("skeleton_black", "skeleton_overlay", "canny")


def _letterbox_scale(size, max_size: int) -> float:
    ow, oh = size
    if max_size <= 0:
        return 1.0
    return min(max_size / float(ow), max_size / float(oh)) if ow and oh else 1.0


def _visible(kp: np.ndarray) -> tuple[bool, ...]:
    return tuple(bool(kp[i, 2] >= _KP_CONF) for i in range(len(kp)))


def _draw_skeleton(canvas: ImageDraw.ImageDraw, kpts: list[np.ndarray], scale: float = 1.0) -> None:
    for pi, kp in enumerate(kpts):
        vis = _visible(kp)
        color = _PALETTE[pi % len(_PALETTE)]
        for a, b in _SKELETON_EDGES:
            if a >= len(kp) or b >= len(kp) or not (vis[a] and vis[b]):
                continue
            ax, ay = kp[a, 0] * scale, kp[a, 1] * scale
            bx, by = kp[b, 0] * scale, kp[b, 1] * scale
            canvas.line([ax, ay, bx, by], fill=color, width=3, joint="curve")
        for i in range(len(kp)):
            if not vis[i]:
                continue
            r = 4 if i <= 4 else 3
            x, y = kp[i, 0] * scale, kp[i, 1] * scale
            canvas.ellipse([x - r, y - r, x + r, y + r], fill=color)


def _skeleton_image(image: Image.Image, kpts: list[np.ndarray], black_bg: bool) -> Image.Image | None:
    if not kpts:
        return None
    rgb = image.convert("RGB")
    if black_bg:
        scale = _letterbox_scale(rgb.size, 640)
        cw, ch = max(1, round(rgb.width * scale)), max(1, round(rgb.height * scale))
        canvas = Image.new("RGB", (cw, ch), (0, 0, 0))
    else:
        scale = 1.0
        canvas = rgb.copy()
    _draw_skeleton(ImageDraw.Draw(canvas), kpts, scale)
    return canvas


def _convolve3(im: np.ndarray, kernel: np.ndarray) -> np.ndarray:
    padded = np.pad(im, 1, mode="reflect")
    out = np.zeros_like(im)
    for i in range(3):
        for j in range(3):
            out += kernel[i, j] * padded[i:i + im.shape[0], j:j + im.shape[1]]
    return out


def _canny_image(image: Image.Image, high: float = 0.35, low: float = 0.10) -> Image.Image:
    gray = np.asarray(image.convert("L"), dtype=np.float64) / 255.0
    blur = _convolve3(gray, np.array([[1, 2, 1], [2, 4, 2], [1, 2, 1]], dtype=np.float64) / 16.0)
    gx = _convolve3(blur, np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=np.float64))
    gy = _convolve3(blur, np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=np.float64))
    mag = np.sqrt(gx * gx + gy * gy)
    peak = float(mag.max())
    if peak > 0:
        mag /= peak
    strong = mag >= high
    weak = (mag >= low) & ~strong
    dilated = strong.copy()
    for dy in (-1, 0, 1):
        for dx in (-1, 0, 1):
            dilated |= np.roll(np.roll(strong, dy, axis=0), dx, axis=1)
    edges = strong | (weak & dilated)
    return Image.fromarray(np.where(edges, 255, 0).astype(np.uint8), mode="L").convert("RGB")


def render_control(image, keypoints: list[np.ndarray] | None, kind: str = "skeleton_black") -> Image.Image | None:
    """Render one ControlNet-style preview.

    ``keypoints`` is a list of COCO-17 arrays in image coordinates (as
    returned by ``PoseEstimator.estimate``). Returns ``None`` when a skeleton
    kind has no keypoints (caller hides the preview); the ``canny`` kind is
    independent of pose.
    """
    pil = image if isinstance(image, Image.Image) else Image.fromarray(np.asarray(image))
    kpts = [np.asarray(k, dtype=np.float64) for k in (keypoints or [])]
    kind = (kind or "skeleton_black").strip().lower()
    if kind == "canny":
        return _canny_image(pil)
    if kind == "skeleton_overlay":
        return _skeleton_image(pil, kpts, black_bg=False)
    return _skeleton_image(pil, kpts, black_bg=True)


def save_control_png(pil: Image.Image) -> str:
    """Persist a preview to a temp file so ``gr.DownloadButton`` can serve it."""
    fd, path = tempfile.mkstemp(prefix="whyx_control_", suffix=".png")
    os.close(fd)
    pil.save(path, "PNG")
    return path