"""Image preprocessing for EKG strips.

For PTB-XL rendered images this is essentially a no-op — the renderer
already produces clean grids on a parchment background. For user-uploaded
phone photos, the pipeline below trims margins, corrects light skew, and
normalizes contrast so the multimodal model sees the cleanest possible
representation.

Each step is conservative: skip if the input looks already-good. The cost
of doing nothing on a clean image is small; the cost of mangling a clean
image with aggressive deskew is large.
"""

from __future__ import annotations

import io
from dataclasses import dataclass
from pathlib import Path

import cv2
import numpy as np
from PIL import Image


@dataclass(frozen=True)
class PreparedImage:
    pil: Image.Image
    width: int
    height: int
    notes: list[str]

    def to_png_bytes(self) -> bytes:
        buf = io.BytesIO()
        self.pil.save(buf, format="PNG", optimize=True)
        return buf.getvalue()


def prepare(path: str | Path) -> PreparedImage:
    """Load an image from disk and return a preprocessed copy.

    The original is never modified. Notes accumulate for any operation
    that was applied so the upstream service can attach them to
    image_quality_notes on the canonical interpretation.
    """
    p = Path(path)
    if not p.exists():
        raise FileNotFoundError(p)

    pil = Image.open(p).convert("RGB")
    notes: list[str] = []

    np_img = np.array(pil)
    bgr = cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR)

    bgr, applied = _autocrop_borders(bgr)
    if applied:
        notes.append("auto-cropped uniform borders")

    bgr, applied = _maybe_deskew(bgr)
    if applied:
        notes.append("rotated to correct skew")

    bgr = _normalize_contrast(bgr)

    rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
    out = Image.fromarray(rgb)
    return PreparedImage(pil=out, width=out.width, height=out.height, notes=notes)


def _autocrop_borders(bgr: np.ndarray) -> tuple[np.ndarray, bool]:
    """Trim solid-color borders. Conservative: bails out if it would crop
    more than 12% in any direction."""
    gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
    # Treat near-white and near-black as "border".
    mask = ((gray > 240) | (gray < 16)).astype(np.uint8) * 255
    inv = cv2.bitwise_not(mask)
    coords = cv2.findNonZero(inv)
    if coords is None:
        return bgr, False
    x, y, w, h = cv2.boundingRect(coords)
    H, W = bgr.shape[:2]
    if x < 4 and y < 4 and (x + w) > W - 4 and (y + h) > H - 4:
        return bgr, False
    # Refuse to crop too aggressively — better to leave the strip intact.
    if w < 0.88 * W or h < 0.88 * H:
        return bgr, False
    return bgr[y : y + h, x : x + w], True


def _maybe_deskew(bgr: np.ndarray) -> tuple[np.ndarray, bool]:
    """Detect dominant near-horizontal lines (EKG grid) and rotate to align.
    Falls back to identity on any ambiguity."""
    gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)
    lines = cv2.HoughLinesP(
        edges, 1, np.pi / 180, threshold=120, minLineLength=120, maxLineGap=10
    )
    if lines is None:
        return bgr, False
    angles: list[float] = []
    for line in lines[:200]:
        x1, y1, x2, y2 = line[0]
        dx = x2 - x1
        dy = y2 - y1
        if dx == 0:
            continue
        ang = np.degrees(np.arctan2(dy, dx))
        if -20 < ang < 20:
            angles.append(ang)
    if len(angles) < 20:
        return bgr, False
    median = float(np.median(angles))
    if abs(median) < 0.5:
        return bgr, False
    H, W = bgr.shape[:2]
    matrix = cv2.getRotationMatrix2D((W / 2, H / 2), median, 1.0)
    rotated = cv2.warpAffine(
        bgr,
        matrix,
        (W, H),
        flags=cv2.INTER_CUBIC,
        borderMode=cv2.BORDER_REPLICATE,
    )
    return rotated, True


def _normalize_contrast(bgr: np.ndarray) -> np.ndarray:
    """Mild CLAHE on the L channel of LAB. Conservative clip limit so the
    parchment-on-grid look stays recognizable."""
    lab = cv2.cvtColor(bgr, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=1.5, tileGridSize=(8, 8))
    l2 = clahe.apply(l)
    return cv2.cvtColor(cv2.merge([l2, a, b]), cv2.COLOR_LAB2BGR)