"""Module for adding text watermarks to videos."""

from __future__ import annotations

import logging
import multiprocessing
import os
from functools import partial
from typing import TYPE_CHECKING, Any, cast

import av
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from tqdm import tqdm

from videolab.utils import generate_output_filename, handle_video_errors

if TYPE_CHECKING:
    import argparse
    from collections.abc import Iterator
    from pathlib import Path

logger = logging.getLogger(__name__)


def _parse_font_color(font_color_str: str) -> tuple[int, int, int, int]:
    """Parse a color string into a validated (R, G, B, A) tuple.

    Supports:
    - RGBA comma-separated string (e.g., "255,255,255,128")
    - RGB hexadecimal string (e.g., "#FFFFFF" or "FFFFFF"), with A defaulted to
      128.
    """
    # 1. check for hexadecimal format
    hex_color = font_color_str.removeprefix("#")

    hex_length = 6
    if len(hex_color) == hex_length and hex_color.isalnum():
        try:
            r = int(hex_color[0:2], 16)
            g = int(hex_color[2:4], 16)
            b = int(hex_color[4:6], 16)
            a = 128  # default alpha for hex codes (50% opacity)
        except ValueError as e:
            msg = (
                f"Invalid hex color format: '{font_color_str}'. "
                "Expected 6 hexadecimal digits."
            )
            raise ValueError(msg) from e
        else:
            return (r, g, b, a)

    # 2. check for RGBA comma-separated format
    try:
        components = list(map(int, font_color_str.split(",")))
    except ValueError as e:
        msg = (
            f"Invalid color format: '{font_color_str}'. "
            "Expected 'R,G,B,A' or 'hex code'."
        )
        raise ValueError(msg) from e

    num_components = 4
    if len(components) != num_components:
        msg = (
            "Comma-separated string must contain exactly 4 components "
            "(R,G,B,A)."
        )
        raise ValueError(msg)

    r, g, b, a = components
    max_color_val = 255
    if not all(0 <= c <= max_color_val for c in components):
        msg = "All color components (R, G, B, A) must be between 0 and 255."
        raise ValueError(msg)

    return (r, g, b, a)


def _calculate_watermark_position(
    position_str: str,
    frame_size: tuple[int, int],
    text_size: tuple[int, int],
    *,
    margin: int,
) -> tuple[int, int]:
    """Calculate the (x, y) position for the watermark."""
    frame_width, frame_height = frame_size
    text_width, text_height = text_size
    if position_str == "top-left":
        x = margin
        y = margin
    elif position_str == "bottom-left":
        x = margin
        y = frame_height - text_height - margin
    elif position_str == "bottom-right":
        x = frame_width - text_width - margin
        y = frame_height - text_height - margin
    elif position_str == "center":
        x = (frame_width - text_width) // 2
        y = (frame_height - text_height) // 2
    else:  # top-right is the default
        x = frame_width - text_width - margin
        y = margin
    return (x, y)


def _create_watermark_image(  # noqa: PLR0913
    text: str,
    frame_size: tuple[int, int],
    font_size_arg: int | None,
    font_color_str: str,
    *,
    position_str: str,
    margin: int,
) -> tuple[np.ndarray, tuple[int, int]]:
    """Create a transparent watermark image with text."""
    _, frame_height = frame_size
    min_font_size = 10
    font_size_factor = 0.05
    font_size = (
        font_size_arg
        if font_size_arg is not None
        else max(min_font_size, int(frame_height * font_size_factor))
    )
    font = ImageFont.load_default(size=font_size)

    default_color = (255, 255, 255, 128)
    try:
        font_color = _parse_font_color(font_color_str)
    except ValueError as e:
        logger.warning(
            "Invalid font color format '%s': %s. Using default %s.",
            font_color_str,
            e,
            default_color,
        )
        font_color = default_color

    # use a dummy image to calculate text size
    dummy_img = Image.new("RGBA", (1, 1))
    draw = ImageDraw.Draw(dummy_img)
    left, top, right, bottom = draw.textbbox((0, 0), text, font=font)
    text_width = int(right - left)
    text_height = int(bottom - top)

    # create the watermark image
    img = Image.new("RGBA", (text_width, text_height), (0, 0, 0, 0))
    draw = ImageDraw.Draw(img)
    draw.text((-left, -top), text, font=font, fill=font_color)

    # convert Pillow Image to NumPy array
    watermark_array = np.array(img)

    position = _calculate_watermark_position(
        position_str,
        frame_size,
        (text_width, text_height),
        margin=margin,
    )

    return watermark_array, position


def _add_watermark_to_array(
    frame_array: np.ndarray,
    watermark_img: np.ndarray,
    position: tuple[int, int],
) -> np.ndarray:
    """Overlay a watermark image (NumPy array) onto a frame array."""
    x, y = position
    h, w, _ = watermark_img.shape

    # extract the alpha channel from the watermark
    alpha_watermark = (
        watermark_img[:, :, 3] / 255.0
    )  # normalized alpha [0.0, 1.0]
    alpha_frame = 1.0 - alpha_watermark

    # apply the watermark using alpha blending (vectorized)
    frame_region = frame_array[y : y + h, x : x + w, 0:3]
    watermark_rgb = watermark_img[:, :, 0:3]

    # the alpha arrays (2D) are broadcasted to the 3rd dimension (color
    # channels) the formula is: C_final = C_frame * (1 - alpha_watermark) +
    # C_watermark * alpha_watermark
    frame_region[:] = (
        alpha_frame[:, :, np.newaxis]
        * frame_region  # background frame component
        + alpha_watermark[:, :, np.newaxis]
        * watermark_rgb  # foreground watermark component
    )

    # ensure the alpha channel of the frame is opaque where watermark is
    # applied
    frame_array[y : y + h, x : x + w, 3] = 255
    return frame_array


def _process_frame_worker(
    frame_data: tuple[np.ndarray, int, av.Rational],
    watermark_img: np.ndarray,
    watermark_position: tuple[int, int],
) -> tuple[np.ndarray, int, av.Rational]:
    """Worker function to process a single frame."""
    frame_array, pts, time_base = frame_data
    processed_array = _add_watermark_to_array(
        frame_array, watermark_img, watermark_position
    )
    return processed_array, pts, time_base


def _add_watermark(
    frame: av.VideoFrame,
    watermark_img: np.ndarray,
    position: tuple[int, int],
) -> av.VideoFrame:
    """Overlay a watermark image (NumPy array) onto a video frame."""
    frame_array = frame.to_ndarray(format="rgba")
    processed_array = _add_watermark_to_array(
        frame_array, watermark_img, position
    )

    new_frame = cast(
        "av.VideoFrame",
        av.VideoFrame.from_ndarray(processed_array, format="rgba"),
    )
    new_frame.pts = frame.pts
    new_frame.time_base = frame.time_base
    return new_frame


def _apply_watermark_to_video(  # noqa: PLR0913
    input_file: str,
    output_file: Path,
    text: str,
    font_size: int | None,
    font_color: str,
    *,
    position: str,
    margin: int,
    num_processes: int | None = None,
) -> None:
    """Core logic to add a text watermark to a video."""
    with av.open(input_file, mode="r") as in_container:
        in_video_stream = next(
            (s for s in in_container.streams if s.type == "video"),
            None,
        )
        if not in_video_stream:
            msg = f"No video stream found in '{input_file}'."
            raise ValueError(msg)

        watermark_img, watermark_position = _create_watermark_image(
            text,
            (in_video_stream.width, in_video_stream.height),
            font_size,
            font_color,
            position_str=position,
            margin=margin,
        )

        with av.open(output_file, mode="w") as out_container:
            out_video_stream = out_container.add_stream_from_template(
                in_video_stream
            )

            total_frames = in_video_stream.frames
            processes = num_processes or os.cpu_count() or 1

            # generator to yield frames for processing
            def frame_generator() -> Iterator[
                tuple[np.ndarray, int, av.Rational]
            ]:
                for packet in in_container.demux(in_video_stream):
                    for frame in packet.decode():
                        # convert to array here because VideoFrame is not
                        # picklable
                        yield (
                            frame.to_ndarray(format="rgba"),
                            frame.pts,
                            frame.time_base,
                        )

            worker_func = partial(
                _process_frame_worker,
                watermark_img=watermark_img,
                watermark_position=watermark_position,
            )

            with (
                multiprocessing.Pool(processes=processes) as pool,
                tqdm(
                    total=total_frames, unit="frame", desc="Watermarking"
                ) as pbar,
            ):
                # imap preserves order, which is critical for video encoding
                # we use a small chunksize to balance IPC overhead and latency
                for processed_array, pts, time_base in pool.imap(
                    worker_func, frame_generator(), chunksize=4
                ):
                    new_frame = av.VideoFrame.from_ndarray(
                        processed_array, format="rgba"
                    )
                    new_frame.pts = pts
                    new_frame.time_base = time_base

                    for new_packet in out_video_stream.encode(new_frame):
                        out_container.mux(new_packet)
                    pbar.update(1)

            # flush the encoder
            for new_packet in out_video_stream.encode():
                out_container.mux(new_packet)


@handle_video_errors
def watermark_command(args: argparse.Namespace) -> None:
    """Add a text watermark to a video."""
    output_file = generate_output_filename(
        args.input_file,
        args.output_file,
        "watermarked",
    )

    _apply_watermark_to_video(
        args.input_file,
        output_file,
        args.text,
        args.font_size,
        args.font_color,
        position=args.position,
        margin=args.margin,
        num_processes=args.cpu,
    )


def register_subcommand(subparsers: argparse._SubParsersAction[Any]) -> None:
    """Register the 'watermark' subcommand."""
    parser = subparsers.add_parser(
        "watermark",
        help="add a watermark to a video.",
        description="Adds a watermark to a video file.",
    )
    parser.add_argument(
        "input_file", type=str, help="Path to the input video file"
    )
    parser.add_argument(
        "output_file",
        type=str,
        nargs="?",
        default=None,
        help=(
            "Path to save the watermarked video file, "
            "defaults: '<input_file>_watermarked.<ext>'"
        ),
    )
    parser.add_argument(
        "--text",
        type=str,
        required=True,
        help="The text for the watermark",
    )
    parser.add_argument(
        "--font-size",
        type=int,
        default=None,
        help=(
            "Font size for the watermark text. "
            "Defaults to 5%% of video height."
        ),
    )
    parser.add_argument(
        "--font-color",
        type=str,
        default="255,255,255,128",
        help="Font color in RGBA format (e.g., '255,255,255,128').",
    )
    parser.add_argument(
        "--position",
        type=str,
        default="top-right",
        choices=[
            "top-left",
            "top-right",
            "bottom-left",
            "bottom-right",
            "center",
        ],
        help="Position of the watermark.",
    )
    parser.add_argument(
        "--margin",
        type=int,
        default=20,
        help="Margin from the edges of the video in pixels.",
    )
    parser.add_argument(
        "--cpu",
        type=int,
        default=None,
        help=(
            "number of CPU cores for parallel processing "
            "(default: all available)"
        ),
    )
    parser.set_defaults(func=watermark_command)
