"""Text overlay module for images and frames.

Provides functionality to add text overlays onto video frames,
supporting multiple pixel formats including YUYV and BGR.
"""

import cv2
import numpy as np

# Default message displayed on publisher frames
PUBLISHER_MESSAGE = "This is a test."


class Color:
    """Color definitions for text overlay.
    
    Colors are defined in BGR format for OpenCV compatibility.
    For YUYV overlay, these are automatically converted.
    """
    
    # BGR colors
    YELLOW_BGR = (0, 255, 255)
    WHITE_BGR = (255, 255, 255)
    RED_BGR = (0, 0, 255)
    GREEN_BGR = (0, 255, 0)
    BLUE_BGR = (255, 0, 0)
    BLACK_BGR = (0, 0, 0)


def bgr_to_yuyv_frame(bgr_frame: np.ndarray) -> np.ndarray:
    """Convert a BGR frame to YUYV format.
    
    Args:
        bgr_frame: Input BGR frame with shape (H, W, 3).
        
    Returns:
        YUYV frame with shape (H, W, 2).
    """
    height, width = bgr_frame.shape[:2]
    
    # Convert BGR to YUV (this gives Y, U, V as separate channels)
    yuv = cv2.cvtColor(bgr_frame, cv2.COLOR_BGR2YUV)
    
    # Extract Y, U, V channels
    y = yuv[:, :, 0]
    u = yuv[:, :, 1]
    v = yuv[:, :, 2]
    
    # Subsample U and V (average pairs of pixels horizontally)
    u_sub = (u[:, ::2].astype(np.uint16) + u[:, 1::2].astype(np.uint16)) // 2
    v_sub = (v[:, ::2].astype(np.uint16) + v[:, 1::2].astype(np.uint16)) // 2
    u_sub = u_sub.astype(np.uint8)
    v_sub = v_sub.astype(np.uint8)
    
    # Create YUYV output: for each pair of pixels (Y0, U, Y1, V)
    # We need output shape (H, W, 2) where each "pixel" is a YUYV pair
    yuyv = np.zeros((height, width, 2), dtype=np.uint8)
    
    # Even columns get Y values
    yuyv[:, ::2, 0] = y[:, ::2]   # Y0
    yuyv[:, 1::2, 0] = y[:, 1::2]  # Y1
    
    # Odd columns get U and V (shared between pairs)
    # U goes at position 0 of odd columns
    yuyv[:, ::2, 1] = u_sub
    # V goes at position 1 of odd columns  
    yuyv[:, 1::2, 1] = v_sub
    
    return yuyv


def convert_yuyv_to_bgr(yuyv_frame: np.ndarray) -> np.ndarray:
    """Convert YUYV frame to BGR format for text rendering.
    
    Args:
        yuyv_frame: Input YUYV frame with shape (H, W, 2).
        
    Returns:
        BGR frame with shape (H, W, 3).
    """
    # Extract Y, U, V channels from YUYV
    height, width = yuyv_frame.shape[:2]
    
    # YUYV is stored as (H, W, 2) where:
    # Channel 0: Y values for all pixels
    # Channel 1: U at even columns, V at odd columns
    y = yuyv_frame[:, :, 0]
    uv = yuyv_frame[:, :, 1]
    
    # Reconstruct full YUV image
    yuv = np.zeros((height, width, 3), dtype=np.uint8)
    yuv[:, :, 0] = y  # Y channel
    
    # U and V are subsampled horizontally, need to upsample
    # U is at even columns, V at odd columns in the UV channel
    u_sub = uv[:, ::2]
    v_sub = uv[:, 1::2]
    
    # Upsample U and V to match Y resolution
    u_full = np.repeat(u_sub, 2, axis=1)[:, :width]
    v_full = np.repeat(v_sub, 2, axis=1)[:, :width]
    
    yuv[:, :, 1] = u_full
    yuv[:, :, 2] = v_full
    
    # Convert YUV to BGR
    bgr = cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR)
    
    return bgr


def bgr_to_yuyv_pixel(b: int, g: int, r: int) -> tuple[int, int]:
    """Convert a single BGR pixel to YUYV (Y and UV values).
    
    Args:
        b: Blue component (0-255).
        g: Green component (0-255).
        r: Red component (0-255).
        
    Returns:
        Tuple of (Y, UV) where UV is the average U/V for the color.
    """
    # BGR to YUV conversion formulas
    y = int(0.299 * r + 0.587 * g + 0.114 * b)
    u = int(-0.169 * r - 0.331 * g + 0.499 * b + 128)
    v = int(0.499 * r - 0.418 * g - 0.0813 * b + 128)
    
    # Clamp to valid range
    y = max(0, min(255, y))
    u = max(0, min(255, u))
    v = max(0, min(255, v))
    
    return y, u, v


def overlay_text_yuyv(
    frame: np.ndarray,
    text: str,
    position: tuple[int, int] = (50, 50),
    color_bgr: tuple[int, int, int] = Color.YELLOW_BGR,
    font_scale: float = 0.7,
    thickness: int = 2,
) -> np.ndarray:
    """Overlay text onto a YUYV frame.
    
    Converts the frame to BGR, renders the text, then converts back to YUYV.
    
    Args:
        frame: Input YUYV frame with shape (H, W, 2).
        text: Text string to overlay.
        position: (x, y) coordinates for text position.
        color_bgr: Text color in BGR format (default: yellow).
        font_scale: Font scale factor.
        thickness: Text thickness in pixels.
        
    Returns:
        YUYV frame with text overlay.
    """
    # Convert YUYV to BGR
    bgr = convert_yuyv_to_bgr(frame)
    
    # Add text to BGR frame
    cv2.putText(
        bgr,
        text,
        position,
        cv2.FONT_HERSHEY_SIMPLEX,
        font_scale,
        color_bgr,
        thickness,
        cv2.LINE_AA,
    )
    
    # Convert back to YUYV
    result = bgr_to_yuyv_frame(bgr)
    
    return result


def overlay_text_bgr(
    frame: np.ndarray,
    text: str,
    position: tuple[int, int] = (50, 50),
    color: tuple[int, int, int] = Color.YELLOW_BGR,
    font_scale: float = 0.7,
    thickness: int = 2,
) -> np.ndarray:
    """Overlay text onto a BGR frame.
    
    Args:
        frame: Input BGR frame with shape (H, W, 3).
        text: Text string to overlay.
        position: (x, y) coordinates for text position.
        color: Text color in BGR format (default: yellow).
        font_scale: Font scale factor.
        thickness: Text thickness in pixels.
        
    Returns:
        BGR frame with text overlay.
    """
    result = frame.copy()
    cv2.putText(
        result,
        text,
        position,
        cv2.FONT_HERSHEY_SIMPLEX,
        font_scale,
        color,
        thickness,
        cv2.LINE_AA,
    )
    return result


def overlay_multiline_text(
    frame: np.ndarray,
    lines: list[str],
    start_position: tuple[int, int] = (50, 50),
    color_bgr: tuple[int, int, int] = Color.YELLOW_BGR,
    font_scale: float = 0.6,
    thickness: int = 1,
    line_spacing: int = 25,
) -> np.ndarray:
    """Overlay multiple lines of text onto a frame.
    
    Args:
        frame: Input frame (YUYV or BGR - auto-detected by shape).
        lines: List of text strings to overlay.
        start_position: (x, y) coordinates for first line.
        color_bgr: Text color in BGR format (default: yellow).
        font_scale: Font scale factor.
        thickness: Text thickness in pixels.
        line_spacing: Vertical spacing between lines.
        
    Returns:
        Frame with text overlay in original format.
    """
    x, y = start_position
    is_yuyv = len(frame.shape) == 3 and frame.shape[2] == 2
    
    if is_yuyv:
        # Convert to BGR for processing
        bgr = convert_yuyv_to_bgr(frame)
    else:
        bgr = frame.copy()
    
    # Add each line
    current_y = y
    for line in lines:
        cv2.putText(
            bgr,
            line,
            (x, current_y),
            cv2.FONT_HERSHEY_SIMPLEX,
            font_scale,
            color_bgr,
            thickness,
            cv2.LINE_AA,
        )
        current_y += line_spacing
    
    # Convert back if needed
    if is_yuyv:
        return bgr_to_yuyv_frame(bgr)
    return bgr


def add_timestamp(
    frame: np.ndarray,
    timestamp: str | None = None,
    position: tuple[int, int] = (10, 30),
    color_bgr: tuple[int, int, int] = Color.WHITE_BGR,
) -> np.ndarray:
    """Add a timestamp to the frame.
    
    Args:
        frame: Input frame (YUYV or BGR - auto-detected).
        timestamp: Timestamp string. If None, current time is used.
        position: (x, y) coordinates.
        color_bgr: Text color in BGR format.
        
    Returns:
        Frame with timestamp in original format.
    """
    if timestamp is None:
        from datetime import datetime
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    return overlay_text_yuyv(
        frame, timestamp, position, color_bgr, font_scale=0.5, thickness=1
    ) if len(frame.shape) == 3 and frame.shape[2] == 2 else overlay_text_bgr(
        frame, timestamp, position, color_bgr, font_scale=0.5, thickness=1
    )


def add_publisher_message(
    frame: np.ndarray,
    message: str = PUBLISHER_MESSAGE,
    position: tuple[int, int] = (50, 50),
    line2: str | None = None,
) -> np.ndarray:
    """Add the publisher message to a frame in yellow.
    
    This is the primary function used by the publisher to add
    the PUBLISHER_MESSAGE to frames.
    
    Args:
        frame: Input frame (YUYV or BGR - auto-detected).
        message: Message text to display (default: PUBLISHER_MESSAGE).
        position: (x, y) coordinates for text.
        line2: Optional second line of text to display below the main message.
        
    Returns:
        Frame with message in original format.
    """
    is_yuyv = len(frame.shape) == 3 and frame.shape[2] == 2
    
    # Get frame dimensions for centering second line
    height, width = frame.shape[:2]
    
    if is_yuyv:
        # Add main message
        result = overlay_text_yuyv(
            frame,
            message,
            position=position,
            color_bgr=Color.YELLOW_BGR,
            font_scale=0.8,
            thickness=2,
        )
        # Add second line if provided
        if line2 is not None:
            # Calculate position for second line (centered, below first line)
            font = cv2.FONT_HERSHEY_SIMPLEX
            font_scale = 0.6
            thickness = 1
            (text_width, text_height), _ = cv2.getTextSize(line2, font, font_scale, thickness)
            x = (width - text_width) // 2
            y = position[1] + 30  # Below first line with spacing
            result = overlay_text_yuyv(
                result,
                line2,
                position=(x, y),
                color_bgr=Color.YELLOW_BGR,
                font_scale=font_scale,
                thickness=thickness,
            )
        return result
    else:
        # Add main message
        result = overlay_text_bgr(
            frame,
            message,
            position=position,
            color=Color.YELLOW_BGR,
            font_scale=0.8,
            thickness=2,
        )
        # Add second line if provided
        if line2 is not None:
            # Calculate position for second line (centered, below first line)
            font = cv2.FONT_HERSHEY_SIMPLEX
            font_scale = 0.6
            thickness = 1
            (text_width, text_height), _ = cv2.getTextSize(line2, font, font_scale, thickness)
            x = (width - text_width) // 2
            y = position[1] + 30  # Below first line with spacing
            result = overlay_text_bgr(
                result,
                line2,
                position=(x, y),
                color=Color.YELLOW_BGR,
                font_scale=font_scale,
                thickness=thickness,
            )
        return result
