import cv2
import numpy as np


def draw_detected_markers(
    image: np.ndarray, corners: list[np.ndarray] | np.ndarray, ids: np.ndarray | None
) -> np.ndarray:
    """
    Draws ArUco markers outlines and IDs on the image.

    Args:
        image: Input image (BGR or grayscale).
        corners: List of marker corners or np.ndarray.
        ids: Array of marker IDs.

    Returns:
        Image with markers drawn.
    """
    if image.ndim == 2:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)

    if corners is not None and len(corners) > 0:
        # Draw outlines
        # corners can be list of (1, 4, 2) or np.ndarray of (N, 4, 2)
        int_corners = [c.astype(np.int32).reshape((-1, 1, 2)) for c in corners]
        cv2.polylines(image, int_corners, True, (0, 255, 0), 2)

        if ids is not None:
            # Flatten ids to handle (N,) or (N, 1)
            flat_ids = ids.flatten()
            for i, corner in enumerate(corners):
                # Put ID near the first corner
                # corner can be (1, 4, 2) or (4, 2)
                c = corner.reshape((-1, 2))
                pos = c[0].astype(int)
                cv2.putText(
                    image,
                    str(flat_ids[i]),
                    tuple(pos),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.5,
                    (0, 0, 255),
                    2,
                )
    return image


def draw_pose_axes(
    image: np.ndarray,
    rvec: np.ndarray,
    tvec: np.ndarray,
    K: np.ndarray,
    dist: np.ndarray | None = None,
    length: float = 0.1,
) -> np.ndarray:
    """
    Draws pose axes on the image.

    Args:
        image: Input image.
        rvec: Rotation vector.
        tvec: Translation vector.
        K: Camera intrinsic matrix.
        dist: Distortion coefficients.
        length: Length of the axes.

    Returns:
        Image with axes drawn.
    """
    if image.ndim == 2:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)

    cv2.drawFrameAxes(image, K, dist, rvec, tvec, length)
    return image


def show_preview(
    frames: dict[int, np.ndarray], *, wait_ms: int = 1, window_prefix: str = "View "
) -> int:
    """
    Shows multiple camera frames in separate windows.

    Args:
        frames: Dictionary mapping camera serials to images.
        wait_ms: Time to wait for a key press in milliseconds.
        window_prefix: Prefix for the window names.

    Returns:
        The key code pressed, or -1 if no key was pressed.
    """
    for serial, frame in frames.items():
        cv2.imshow(f"{window_prefix}{serial}", frame)
    return cv2.waitKey(wait_ms)