from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path

import cv2
import numpy as np

PALETTE: tuple[tuple[int, int, int], ...] = (
    (60, 60, 255),
    (60, 220, 255),
    (90, 230, 90),
    (255, 190, 60),
    (255, 110, 200),
    (255, 255, 255),
)

COLOR_NAMES: tuple[str, ...] = (
    "red", "yellow", "green", "blue", "pink", "white",
)


@dataclass
class Stroke:
    """One drawn stroke."""

    color: tuple[int, int, int]
    thickness: int
    points: list[tuple[int, int]] = field(default_factory=list)


class Canvas:
    """Vector drawing canvas."""

    MAX_JUMP = 0.12

    def __init__(self, width: int, height: int, output_dir: str | Path = "output") -> None:
        self.width = width
        self.height = height
        self.output_dir = Path(output_dir)
        self.strokes: list[Stroke] = []
        self._active: Stroke | None = None
        self.color: tuple[int, int, int] = PALETTE[0]
        self.thickness: int = 6
        self._layer = np.zeros((height, width, 3), dtype=np.uint8)
        self._mask = np.zeros((height, width), dtype=np.uint8)
        self._bounds: list[int] | None = None
        self.revision = 0

    def begin(self) -> None:
        """Start a new stroke."""
        if self._active is None:
            self._active = Stroke(color=self.color, thickness=self.thickness)
            self.strokes.append(self._active)

    def add_point(self, point) -> None:
        """Append point to stroke."""
        if self._active is None:
            self.begin()
        assert self._active is not None
        x, y = int(round(float(point[0]))), int(round(float(point[1])))
        x = max(0, min(self.width - 1, x))
        y = max(0, min(self.height - 1, y))
        pts = self._active.points
        if pts and pts[-1] == (x, y):
            return
        if pts:
            jump = np.hypot(x - pts[-1][0], y - pts[-1][1])
            if jump > self.MAX_JUMP * np.hypot(self.width, self.height):
                self.end()
                self.begin()
                pts = self._active.points
        pts.append((x, y))
        if len(pts) >= 2:
            self._draw_segment(pts[-2], pts[-1], self._active)
        else:
            self._draw_dot((x, y), self._active)

    def end(self) -> None:
        """Finish current stroke."""
        if self._active is not None and len(self._active.points) < 2:
            if len(self._active.points) == 0:
                self.strokes.remove(self._active)
                self._rebuild()
        self._active = None

    def _grow_bounds(self, points, radius: int) -> None:
        """Extend content bounding box."""
        xs = [p[0] for p in points]
        ys = [p[1] for p in points]
        x0, y0 = min(xs) - radius, min(ys) - radius
        x1, y1 = max(xs) + radius, max(ys) + radius
        if self._bounds is None:
            self._bounds = [x0, y0, x1, y1]
        else:
            b = self._bounds
            b[0], b[1] = min(b[0], x0), min(b[1], y0)
            b[2], b[3] = max(b[2], x1), max(b[3], y1)

    def _draw_segment(self, a, b, stroke: Stroke) -> None:
        """Draw one line segment."""
        cv2.line(self._layer, a, b, stroke.color, stroke.thickness, cv2.LINE_AA)
        cv2.line(self._mask, a, b, 255, stroke.thickness, cv2.LINE_AA)
        self._grow_bounds((a, b), stroke.thickness // 2 + 2)
        self.revision += 1

    def _draw_dot(self, p, stroke: Stroke) -> None:
        """Draw one dot."""
        r = max(1, stroke.thickness // 2)
        cv2.circle(self._layer, p, r, stroke.color, -1, cv2.LINE_AA)
        cv2.circle(self._mask, p, r, 255, -1, cv2.LINE_AA)
        self._grow_bounds((p,), r + 2)
        self.revision += 1

    def _rebuild(self) -> None:
        """Redraw all strokes."""
        self._layer[:] = 0
        self._mask[:] = 0
        self._bounds = None
        self.revision += 1
        for s in self.strokes:
            if len(s.points) == 1:
                self._draw_dot(s.points[0], s)
            for a, b in zip(s.points, s.points[1:]):
                self._draw_segment(a, b, s)

    def undo(self) -> None:
        """Remove last stroke."""
        if self.strokes:
            self.strokes.pop()
            self._active = None
            self._rebuild()

    def clear(self) -> None:
        """Erase whole canvas."""
        self.strokes.clear()
        self._active = None
        self._rebuild()

    def scale_content(self, factor: float) -> None:
        """Scale drawing about center."""
        if not self.strokes or abs(factor - 1.0) < 1e-3:
            return
        cx, cy = self.width * 0.5, self.height * 0.5
        for s in self.strokes:
            s.points = [
                (
                    int(round(cx + (x - cx) * factor)),
                    int(round(cy + (y - cy) * factor)),
                )
                for x, y in s.points
            ]
            s.thickness = int(max(1, min(48, round(s.thickness * factor))))
        self._active = None
        self._rebuild()

    def content_span(self) -> float:
        """Drawing diagonal length."""
        box = self.content_bbox(padding=0)
        if box is None:
            return 0.0
        x0, y0, x1, y1 = box
        return float(np.hypot(x1 - x0, y1 - y0))

    def set_color(self, index: int) -> None:
        """Select palette color."""
        self.color = PALETTE[index % len(PALETTE)]

    def set_thickness(self, value: int) -> None:
        """Set stroke thickness."""
        self.thickness = int(max(1, min(48, value)))

    @property
    def is_empty(self) -> bool:
        """No strokes drawn."""
        return not any(s.points for s in self.strokes)

    @property
    def layer(self) -> np.ndarray:
        """Color drawing layer."""
        return self._layer

    @property
    def mask(self) -> np.ndarray:
        """Drawing alpha mask."""
        return self._mask

    def composite_over(self, frame: np.ndarray, opacity: float = 1.0) -> np.ndarray:
        """Blend drawing onto frame."""
        out = frame.copy()
        box = self.content_bbox(padding=2)
        if box is None:
            return out
        x0, y0, x1, y1 = box
        roi = out[y0:y1, x0:x1]
        layer = self._layer[y0:y1, x0:x1]
        mask = self._mask[y0:y1, x0:x1]
        if opacity < 1.0:
            mask = (mask.astype(np.float32) * opacity).astype(np.uint8)
        a = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
        inv = cv2.bitwise_not(a)
        blended = cv2.add(
            cv2.multiply(roi, inv, scale=1 / 255.0),
            cv2.multiply(layer, a, scale=1 / 255.0),
        )
        out[y0:y1, x0:x1] = blended
        return out

    def to_bgra(self) -> np.ndarray:
        """Drawing with transparent background."""
        bgra = np.dstack([self._layer, self._mask])
        return bgra

    def content_bbox(self, padding: int = 24) -> tuple[int, int, int, int] | None:
        """Bounding box of drawing."""
        if self._bounds is None:
            return None
        bx0, by0, bx1, by1 = self._bounds
        x0 = max(0, bx0 - padding)
        y0 = max(0, by0 - padding)
        x1 = min(self.width, bx1 + padding + 1)
        y1 = min(self.height, by1 + padding + 1)
        if x1 <= x0 or y1 <= y0:
            return None
        return x0, y0, x1, y1

    def cropped_bgra(self, padding: int = 24) -> np.ndarray | None:
        """Cropped transparent drawing."""
        box = self.content_bbox(padding)
        if box is None:
            return None
        x0, y0, x1, y1 = box
        return self.to_bgra()[y0:y1, x0:x1]

    def save(self, tag: str = "") -> list[Path]:
        """Save canvas as PNGs."""
        self.output_dir.mkdir(parents=True, exist_ok=True)
        stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        suffix = f"_{tag}" if tag else ""
        saved: list[Path] = []

        transparent = self.output_dir / f"drawing_{stamp}{suffix}.png"
        cv2.imwrite(str(transparent), self.to_bgra())
        saved.append(transparent)

        white = np.full((self.height, self.width, 3), 255, dtype=np.uint8)
        alpha = (self._mask.astype(np.float32) / 255.0)[:, :, None]
        flat = (white * (1.0 - alpha) + self._layer * alpha).astype(np.uint8)
        on_white = self.output_dir / f"drawing_{stamp}{suffix}_white.png"
        cv2.imwrite(str(on_white), flat)
        saved.append(on_white)

        return saved