from __future__ import annotations

import base64
import binascii
import math
from io import BytesIO

from PIL import Image, ImageDraw

# Constants
_DATA_URL_MARKER = ';base64,'
_DATA_URL_PREFIX = 'data:'
_DEFAULT_JPEG_QUALITY = 75
_DEFAULT_COMPRESSION_QUALITY = 40
_MIN_COMPRESSION_QUALITY = 10
_MAX_COMPRESSION_QUALITY = 95
_DEFAULT_DOWNSCALE_RATIO = 0.85
_QUALITY_REDUCTION_FACTOR = 0.75
_MIN_IMAGE_DIMENSION = 240
_DEFAULT_POINT_RADIUS = 30
_DEFAULT_ARROW_LINE_WIDTH = 10
_DEFAULT_ARROW_HEAD_LENGTH = 100
_DEFAULT_ARROW_HEAD_ANGLE = 30.0
_WHITE_RGB = (255, 255, 255)


def _split_data_url(data: str) -> tuple[str | None, str]:
    """Split a data URL into MIME type and base64 payload."""
    stripped = data.strip()
    if stripped.startswith(_DATA_URL_PREFIX) and _DATA_URL_MARKER in stripped:
        header, payload = stripped.split(',', 1)
        mime = header[len(_DATA_URL_PREFIX) :].split(';', 1)[0]
        return mime, payload
    return None, stripped


def _normalize_format(fmt: str | None) -> str | None:
    """Normalize image format string."""
    if not fmt:
        return None
    upper = fmt.upper()
    return 'JPEG' if upper == 'JPG' else upper


def _decode_base64_image(base64_image: str) -> tuple[Image.Image, str | None]:
    """Decode base64 image string to PIL Image object."""
    mime, payload = _split_data_url(base64_image)
    try:
        image_bytes = base64.b64decode(payload, validate=False)
    except (ValueError, binascii.Error) as exc:
        msg = f'Invalid base64 image payload: {exc}'
        raise ValueError(msg) from exc

    try:
        image = Image.open(BytesIO(image_bytes))
        image.load()
    except Exception as exc:
        msg = f'Invalid image data in base64 payload: {exc}'
        raise ValueError(msg) from exc

    fmt = image.format
    if fmt is None and mime:
        fmt = mime.split('/')[-1]

    return image, _normalize_format(fmt)


def _prepare_for_jpeg(image: Image.Image) -> Image.Image:
    """Prepare image for JPEG encoding by handling transparency."""
    if image.mode in ('RGBA', 'LA'):
        background = Image.new('RGB', image.size, _WHITE_RGB)
        background.paste(image, mask=image.split()[-1])
        return background
    if image.mode != 'RGB':
        return image.convert('RGB')
    return image


def _encode_image_bytes(
    image: Image.Image,
    *,
    fmt: str | None,
    quality: int | None = None,
) -> bytes:
    """Encode PIL Image to bytes with specified format and quality."""
    output = BytesIO()
    use_format = _normalize_format(fmt) or 'PNG'
    save_kwargs: dict[str, object] = {}

    if use_format == 'JPEG':
        image = _prepare_for_jpeg(image)
        save_kwargs['quality'] = quality if quality is not None else _DEFAULT_JPEG_QUALITY
        save_kwargs['optimize'] = True
        save_kwargs['progressive'] = True
    elif use_format == 'PNG':
        save_kwargs['optimize'] = True

    image.save(output, format=use_format, **save_kwargs)
    return output.getvalue()


def _encode_image_to_base64(
    image: Image.Image,
    *,
    fmt: str | None,
    quality: int | None = None,
) -> str:
    """Encode PIL Image to base64 string."""
    return base64.b64encode(_encode_image_bytes(image, fmt=fmt, quality=quality)).decode('ascii')


def draw_red_point_base64(base64_image: str, x: int, y: int, radius: int = _DEFAULT_POINT_RADIUS) -> str:
    """Draw a red point on the image at the specified coordinates."""
    image, fmt = _decode_base64_image(base64_image)
    draw = ImageDraw.Draw(image)
    bbox = (x - radius, y - radius, x + radius, y + radius)
    draw.ellipse(bbox, fill='red', outline='red')
    return _encode_image_to_base64(image, fmt=fmt)


def draw_red_arrow_base64(
    base64_image: str,
    start_x: int,
    start_y: int,
    end_x: int,
    end_y: int,
    *,
    line_width: int = _DEFAULT_ARROW_LINE_WIDTH,
    head_length: int = _DEFAULT_ARROW_HEAD_LENGTH,
    head_angle_deg: float = _DEFAULT_ARROW_HEAD_ANGLE,
) -> str:
    """Draw a red arrow on the image from start to end coordinates."""
    image, fmt = _decode_base64_image(base64_image)
    draw = ImageDraw.Draw(image)

    # Draw arrow shaft
    draw.line([(start_x, start_y), (end_x, end_y)], fill='red', width=line_width)

    # Calculate arrowhead
    angle = math.atan2(end_y - start_y, end_x - start_x)
    head_angle = math.radians(head_angle_deg)

    left_x = end_x - head_length * math.cos(angle - head_angle)
    left_y = end_y - head_length * math.sin(angle - head_angle)
    right_x = end_x - head_length * math.cos(angle + head_angle)
    right_y = end_y - head_length * math.sin(angle + head_angle)

    # Draw arrowhead
    draw.polygon([(end_x, end_y), (left_x, left_y), (right_x, right_y)], fill='red')

    return _encode_image_to_base64(image, fmt=fmt)


def compress_base64_image(
    base64_image: str,
    *,
    max_width: int | None = None,
    max_height: int | None = None,
    quality: int = _DEFAULT_COMPRESSION_QUALITY,
    output_format: str | None = None,
    target_max_bytes: int | None = None,
    min_quality: int = _MIN_COMPRESSION_QUALITY,
    downscale_ratio: float = _DEFAULT_DOWNSCALE_RATIO,
    min_width: int = _MIN_IMAGE_DIMENSION,
    min_height: int = _MIN_IMAGE_DIMENSION,
) -> str | None:
    """Compress a base64 image with optional size and quality constraints.

    Args:
        base64_image: Input image as base64 string
        max_width: Maximum width constraint
        max_height: Maximum height constraint
        quality: Initial JPEG quality (0-100)
        output_format: Output format (default: JPEG)
        target_max_bytes: Target maximum file size in bytes
        min_quality: Minimum quality threshold
        downscale_ratio: Ratio to downscale when quality adjustment isn't enough
        min_width: Minimum width threshold
        min_height: Minimum height threshold

    Returns:
        Compressed image as base64 string, or None if compression fails
    """
    image, fmt = _decode_base64_image(base64_image)

    output_format = output_format or 'JPEG'

    # Resize if dimensions exceed max constraints
    if max_width is not None or max_height is not None:
        target_width = max_width if max_width is not None else image.width
        target_height = max_height if max_height is not None else image.height
        if image.width > target_width or image.height > target_height:
            image.thumbnail((target_width, target_height), Image.Resampling.LANCZOS)

    # If no byte target, encode once and return
    if target_max_bytes is None:
        return _encode_image_to_base64(image, fmt=output_format or fmt, quality=quality)

    # Iteratively reduce quality and size to meet byte target
    current_quality = max(min_quality, min(quality, _MAX_COMPRESSION_QUALITY))
    working = image

    while True:
        encoded = _encode_image_bytes(working, fmt=output_format or fmt, quality=current_quality)
        if len(encoded) <= target_max_bytes:
            return base64.b64encode(encoded).decode('ascii')

        # Try reducing quality first
        if current_quality > min_quality:
            current_quality = max(min_quality, int(current_quality * _QUALITY_REDUCTION_FACTOR))
            continue

        # Quality at minimum, try downscaling
        next_width = max(min_width, int(working.width * downscale_ratio))
        next_height = max(min_height, int(working.height * downscale_ratio))

        # Can't downscale further, return current result
        if next_width == working.width and next_height == working.height:
            return base64.b64encode(encoded).decode('ascii')

        working = working.copy()
        working.thumbnail((next_width, next_height), Image.Resampling.LANCZOS)
        current_quality = max(min_quality, min(quality, _MAX_COMPRESSION_QUALITY))
