from __future__ import annotations

import argparse
import csv
import math
import shutil
import subprocess
from dataclasses import dataclass
from pathlib import Path
from tempfile import TemporaryDirectory

import numpy as np
from PIL import Image, ImageOps


@dataclass(frozen=True)
class Canvas:
    width: int = 588
    height: int = 495
    cx: float = 304.0
    cy: float = 256.0


@dataclass(frozen=True)
class Style:
    background: str = "#3d3d3d"
    material: str = "#eeeeee"
    outline: str = "#f4f4f4"
    outline_width: float = 3.0


@dataclass(frozen=True)
class Geometry:
    root_radius: float = 213.0
    tooth_depth: float = 17.5
    reference_teeth: int = 60
    pitch_radius_offset: float = 8.0
    tooth_sharpness: float = 1.32
    low_tooth_threshold: int = 25
    low_tooth_half_width: float = 0.30
    low_tooth_cusp: float = 0.55
    central_hole_radius: float = 46.0
    bolt_hole_radius: float = 16.5
    material_margin: float = 8.0
    bolt_hole_min_teeth: int = 25
    decorative_cutout_min_scale: float = 0.62
    samples_per_tooth: int = 28


@dataclass(frozen=True)
class SprocketSizing:
    pitch_radius: float
    calculated_root_radius: float
    root_radius: float
    tooth_depth: float
    outer_radius: float
    decorative_scale: float
    bolt_scale: float
    include_bolt_holes: bool
    include_decorative_cutouts: bool


@dataclass(frozen=True)
class Cutout:
    kind: str
    cx: float
    cy: float
    rx: float
    ry: float | None = None
    rotation: float = 0.0


CANVAS = Canvas()
STYLE = Style()
GEOMETRY = Geometry()


def fmt(value: float) -> str:
    return f"{value:.3f}".rstrip("0").rstrip(".")


def polygon_path(points: list[tuple[float, float]]) -> str:
    start = points[0]
    rest = points[1:]
    commands = [f"M {fmt(start[0])} {fmt(start[1])}"]
    commands.extend(f"L {fmt(x)} {fmt(y)}" for x, y in rest)
    commands.append("Z")
    return " ".join(commands)


def circle_points(cx: float, cy: float, radius: float, samples: int = 144) -> list[tuple[float, float]]:
    return [
        (cx + math.cos(t) * radius, cy + math.sin(t) * radius)
        for t in np.linspace(0.0, 2.0 * math.pi, samples, endpoint=False)
    ]


def ellipse_points(
    cx: float,
    cy: float,
    rx: float,
    ry: float,
    rotation_deg: float = 0.0,
    samples: int = 192,
) -> list[tuple[float, float]]:
    rotation = math.radians(rotation_deg)
    cos_r = math.cos(rotation)
    sin_r = math.sin(rotation)
    points = []
    for t in np.linspace(0.0, 2.0 * math.pi, samples, endpoint=False):
        x = math.cos(t) * rx
        y = math.sin(t) * ry
        points.append((cx + x * cos_r - y * sin_r, cy + x * sin_r + y * cos_r))
    return points


DECORATIVE_CUTOUT_SPECS = (
    ("ellipse", 0.0, -141.0, 121.0, 41.0, 0.0),
    ("ellipse", -109.0, 58.0, 106.0, 41.0, 63.0),
    ("ellipse", 109.0, 58.0, 106.0, 41.0, -63.0),
)

BOLT_HOLE_OFFSETS = (
    (-134.0, -76.0),
    (134.0, -76.0),
    (0.0, 157.0),
)

REFERENCE_BOLT_OFFSET_RADIUS = max(math.hypot(dx, dy) for dx, dy in BOLT_HOLE_OFFSETS)


def decorative_reference_extent() -> float:
    extent = 0.0
    for _, dx, dy, rx, ry, rotation in DECORATIVE_CUTOUT_SPECS:
        points = ellipse_points(dx, dy, rx, ry, rotation, 240)
        extent = max(extent, *(math.hypot(x, y) for x, y in points))
    return extent


DECORATIVE_REFERENCE_EXTENT = decorative_reference_extent()


def clamp(value: float, minimum: float, maximum: float) -> float:
    return max(minimum, min(maximum, value))


def calculated_pitch_radius(teeth: int, geometry: Geometry = GEOMETRY) -> float:
    reference_pitch_radius = geometry.root_radius + geometry.pitch_radius_offset
    return reference_pitch_radius * math.sin(math.pi / geometry.reference_teeth) / math.sin(math.pi / teeth)


def sprocket_sizing(teeth: int, geometry: Geometry = GEOMETRY) -> SprocketSizing:
    pitch_radius = calculated_pitch_radius(teeth, geometry)
    calculated_root_radius = pitch_radius - geometry.pitch_radius_offset
    central_min_root_radius = geometry.central_hole_radius + geometry.material_margin

    include_bolt_holes = teeth >= geometry.bolt_hole_min_teeth
    if include_bolt_holes:
        bolt_min_radius = geometry.central_hole_radius + geometry.bolt_hole_radius + geometry.material_margin
        bolt_min_root_radius = bolt_min_radius + geometry.bolt_hole_radius + geometry.material_margin
        min_root_radius = max(central_min_root_radius, bolt_min_root_radius)
    else:
        min_root_radius = central_min_root_radius

    root_radius = max(calculated_root_radius, min_root_radius)
    outer_radius = root_radius + geometry.tooth_depth

    max_bolt_offset = max(0.0, root_radius - geometry.bolt_hole_radius - geometry.material_margin)
    if include_bolt_holes:
        min_bolt_offset = geometry.central_hole_radius + geometry.bolt_hole_radius + geometry.material_margin
        bolt_offset = clamp(max_bolt_offset, min_bolt_offset, REFERENCE_BOLT_OFFSET_RADIUS)
        bolt_scale = bolt_offset / REFERENCE_BOLT_OFFSET_RADIUS
    else:
        bolt_scale = 0.0

    decorative_scale = clamp((root_radius - geometry.material_margin) / DECORATIVE_REFERENCE_EXTENT, 0.0, 1.0)
    include_decorative_cutouts = decorative_scale >= geometry.decorative_cutout_min_scale

    return SprocketSizing(
        pitch_radius=pitch_radius,
        calculated_root_radius=calculated_root_radius,
        root_radius=root_radius,
        tooth_depth=geometry.tooth_depth,
        outer_radius=outer_radius,
        decorative_scale=decorative_scale,
        bolt_scale=bolt_scale,
        include_bolt_holes=include_bolt_holes,
        include_decorative_cutouts=include_decorative_cutouts,
    )


def outer_points(
    teeth: int,
    geometry: Geometry = GEOMETRY,
    canvas: Canvas = CANVAS,
    radius_delta: float = 0.0,
) -> list[tuple[float, float]]:
    samples = teeth * geometry.samples_per_tooth
    sizing = sprocket_sizing(teeth, geometry)
    points = []
    phase_offset = -math.pi / 2.0

    for index in range(samples):
        theta = phase_offset + 2.0 * math.pi * index / samples
        tooth_phase = (index % geometry.samples_per_tooth) / geometry.samples_per_tooth
        radius = sizing.root_radius + sizing.tooth_depth * tooth_profile(teeth, tooth_phase, geometry) + radius_delta
        points.append((canvas.cx + math.cos(theta) * radius, canvas.cy + math.sin(theta) * radius))

    return points


def tooth_profile(teeth: int, phase: float, geometry: Geometry = GEOMETRY) -> float:
    if teeth < geometry.low_tooth_threshold:
        distance_from_peak = abs(phase - 0.5)
        if distance_from_peak >= geometry.low_tooth_half_width:
            return 0.0
        return 1.0 - (distance_from_peak / geometry.low_tooth_half_width) ** geometry.low_tooth_cusp

    smooth_profile = 0.5 - 0.5 * math.cos(2.0 * math.pi * phase)
    return smooth_profile**geometry.tooth_sharpness


def offset_cutout(
    kind: str,
    dx: float,
    dy: float,
    rx: float,
    ry: float | None,
    rotation: float,
    center_scale: float,
    radius_scale: float,
    canvas: Canvas = CANVAS,
) -> Cutout:
    scaled_ry = None if ry is None else ry * radius_scale
    return Cutout(
        kind,
        canvas.cx + dx * center_scale,
        canvas.cy + dy * center_scale,
        rx * radius_scale,
        scaled_ry,
        rotation,
    )


def cutouts_for_teeth(teeth: int, canvas: Canvas = CANVAS, geometry: Geometry = GEOMETRY) -> list[Cutout]:
    sizing = sprocket_sizing(teeth, geometry)
    return [
        *(
            offset_cutout(kind, dx, dy, rx, ry, rotation, sizing.decorative_scale, sizing.decorative_scale, canvas)
            for kind, dx, dy, rx, ry, rotation in DECORATIVE_CUTOUT_SPECS
            if sizing.include_decorative_cutouts
        ),
        offset_cutout("circle", 0.0, 0.0, geometry.central_hole_radius, None, 0.0, 1.0, 1.0, canvas),
        *(
            offset_cutout("circle", dx, dy, geometry.bolt_hole_radius, None, 0.0, sizing.bolt_scale, 1.0, canvas)
            for dx, dy in BOLT_HOLE_OFFSETS
            if sizing.include_bolt_holes
        ),
    ]


def cutout_path(cutout: Cutout, delta: float = 0.0, reverse: bool = False) -> str:
    if cutout.kind == "circle":
        points = circle_points(cutout.cx, cutout.cy, cutout.rx + delta, 160)
    elif cutout.kind == "ellipse":
        assert cutout.ry is not None
        points = ellipse_points(
            cutout.cx,
            cutout.cy,
            cutout.rx + delta,
            cutout.ry + delta,
            cutout.rotation,
            240,
        )
    else:
        raise ValueError(f"Unsupported cutout kind: {cutout.kind}")

    if reverse:
        points = list(reversed(points))
    return polygon_path(points)


def cutout_paths_for_teeth(teeth: int, canvas: Canvas = CANVAS) -> list[str]:
    return [cutout_path(cutout) for cutout in cutouts_for_teeth(teeth, canvas)]


def filled_svg(teeth: int, canvas: Canvas = CANVAS, style: Style = STYLE) -> str:
    outer = polygon_path(outer_points(teeth))
    cutouts = " ".join(cutout_paths_for_teeth(teeth, canvas))
    return f"""<svg xmlns="http://www.w3.org/2000/svg" width="{canvas.width}" height="{canvas.height}" viewBox="0 0 {canvas.width} {canvas.height}">
  <path d="{outer} {cutouts}" fill="{style.material}" fill-rule="evenodd"/>
</svg>
"""


def outline_svg(teeth: int, canvas: Canvas = CANVAS, style: Style = STYLE) -> str:
    half_width = style.outline_width / 2.0
    outer = polygon_path(outer_points(teeth, radius_delta=half_width))
    inner = polygon_path(list(reversed(outer_points(teeth, radius_delta=-half_width))))
    paths = [f"{outer} {inner}"]
    for cutout in cutouts_for_teeth(teeth, canvas):
        paths.append(f"{cutout_path(cutout, half_width)} {cutout_path(cutout, -half_width, reverse=True)}")

    path_elements = "\n  ".join(
        f'<path d="{path}" fill="{style.outline}" fill-rule="evenodd"/>'
        for path in paths
    )
    return f"""<svg xmlns="http://www.w3.org/2000/svg" width="{canvas.width}" height="{canvas.height}" viewBox="0 0 {canvas.width} {canvas.height}">
  {path_elements}
</svg>
"""


def render_png(svg_path: Path, png_path: Path) -> None:
    magick = shutil.which("magick") or shutil.which("convert")
    if magick is None:
        raise RuntimeError("ImageMagick is required for SVG-to-PNG conversion, but neither 'magick' nor 'convert' is on PATH.")

    png_path.parent.mkdir(parents=True, exist_ok=True)
    command = [
        magick,
        "-background",
        "none",
        "-density",
        "192",
        str(svg_path),
        "-resize",
        f"{CANVAS.width}x{CANVAS.height}!",
        str(png_path),
    ]
    subprocess.run(command, check=True)


def write_asset(teeth: int, variant: str, svg_root: Path, png_root: Path, png: bool = True) -> tuple[Path, Path | None]:
    if variant == "filled":
        svg_text = filled_svg(teeth)
    elif variant == "outline":
        svg_text = outline_svg(teeth)
    else:
        raise ValueError(f"Unsupported variant: {variant}")

    stem = f"chainring_{teeth:02d}_{variant}"
    svg_path = svg_root / variant / f"{stem}.svg"
    png_path = png_root / variant / f"{stem}.png"
    svg_path.parent.mkdir(parents=True, exist_ok=True)
    svg_path.write_text(svg_text, encoding="utf-8")

    if png:
        render_png(svg_path, png_path)
        return svg_path, png_path

    return svg_path, None


def generate_all(min_teeth: int, max_teeth: int, output_dir: Path, png: bool = True) -> list[Path]:
    svg_root = output_dir / "svg"
    png_root = output_dir / "png"
    png_paths: list[Path] = []

    for teeth in range(min_teeth, max_teeth + 1):
        for variant in ("filled", "outline"):
            _, png_path = write_asset(teeth, variant, svg_root, png_root, png=png)
            if png_path is not None:
                png_paths.append(png_path)

    return png_paths


def write_radii_manifest(min_teeth: int, max_teeth: int, output_path: Path) -> None:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    fields = [
        "teeth",
        "pitch_radius_px",
        "calculated_root_radius_px",
        "root_radius_px",
        "tooth_depth_px",
        "outer_radius_px",
        "bolt_scale",
        "decorative_scale",
        "bolt_holes",
        "decorative_cutouts",
    ]
    with output_path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        for teeth in range(min_teeth, max_teeth + 1):
            sizing = sprocket_sizing(teeth)
            writer.writerow(
                {
                    "teeth": teeth,
                    "pitch_radius_px": f"{sizing.pitch_radius:.3f}",
                    "calculated_root_radius_px": f"{sizing.calculated_root_radius:.3f}",
                    "root_radius_px": f"{sizing.root_radius:.3f}",
                    "tooth_depth_px": f"{sizing.tooth_depth:.3f}",
                    "outer_radius_px": f"{sizing.outer_radius:.3f}",
                    "bolt_scale": f"{sizing.bolt_scale:.3f}",
                    "decorative_scale": f"{sizing.decorative_scale:.3f}",
                    "bolt_holes": int(sizing.include_bolt_holes),
                    "decorative_cutouts": int(sizing.include_decorative_cutouts),
                }
            )


def make_contact_sheet(png_paths: list[Path], output_path: Path, columns: int = 6) -> None:
    if not png_paths:
        return

    thumbs = []
    thumb_size = (196, 165)
    for path in png_paths:
        image = Image.open(path).convert("RGBA")
        image.thumbnail(thumb_size, Image.Resampling.LANCZOS)
        tile = Image.new("RGBA", thumb_size, (61, 61, 61, 255))
        x = (thumb_size[0] - image.width) // 2
        y = (thumb_size[1] - image.height) // 2
        tile.alpha_composite(image, (x, y))
        thumbs.append(tile)

    rows = math.ceil(len(thumbs) / columns)
    sheet = Image.new("RGBA", (columns * thumb_size[0], rows * thumb_size[1]), (36, 36, 36, 255))
    for index, thumb in enumerate(thumbs):
        x = (index % columns) * thumb_size[0]
        y = (index // columns) * thumb_size[1]
        sheet.alpha_composite(thumb, (x, y))

    output_path.parent.mkdir(parents=True, exist_ok=True)
    sheet.save(output_path)


def compare_to_reference(candidate_path: Path, reference_path: Path) -> dict[str, float]:
    candidate = Image.open(candidate_path).convert("L")
    reference = Image.open(reference_path).convert("L").resize(candidate.size, Image.Resampling.LANCZOS)

    candidate_arr = np.asarray(ImageOps.autocontrast(candidate), dtype=np.float32) / 255.0
    reference_arr = np.asarray(ImageOps.autocontrast(reference), dtype=np.float32) / 255.0
    diff = candidate_arr - reference_arr

    candidate_edges = np.abs(np.gradient(candidate_arr)[0]) + np.abs(np.gradient(candidate_arr)[1])
    reference_edges = np.abs(np.gradient(reference_arr)[0]) + np.abs(np.gradient(reference_arr)[1])
    edge_diff = candidate_edges - reference_edges

    return {
        "mse": float(np.mean(diff**2)),
        "mae": float(np.mean(np.abs(diff))),
        "edge_mae": float(np.mean(np.abs(edge_diff))),
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate chainring SVG and PNG assets with fixed reference-like inner cutouts.")
    parser.add_argument("--min-teeth", type=int, default=5)
    parser.add_argument("--max-teeth", type=int, default=60)
    parser.add_argument("--output", type=Path, default=Path("out"))
    parser.add_argument("--no-png", action="store_true")
    parser.add_argument("--contact-sheet", action="store_true")
    parser.add_argument("--reference", type=Path, help="Optional local reference image for reporting similarity metrics.")
    parser.add_argument("--reference-teeth", type=int, default=44)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.min_teeth < 3:
        raise ValueError("--min-teeth must be at least 3.")
    if args.max_teeth < args.min_teeth:
        raise ValueError("--max-teeth must be greater than or equal to --min-teeth.")

    png_paths = generate_all(args.min_teeth, args.max_teeth, args.output, png=not args.no_png)
    write_radii_manifest(args.min_teeth, args.max_teeth, args.output / "radii.csv")

    if args.contact_sheet and png_paths:
        filled_samples = [path for path in png_paths if path.parent.name == "filled"]
        outline_samples = [path for path in png_paths if path.parent.name == "outline"]
        make_contact_sheet(filled_samples, args.output / "contact_sheet_filled.png")
        make_contact_sheet(outline_samples, args.output / "contact_sheet_outline.png")

    if args.reference and not args.no_png:
        with TemporaryDirectory() as tmp:
            tmp_root = Path(tmp)
            _, candidate = write_asset(args.reference_teeth, "filled", tmp_root / "svg", tmp_root / "png", png=True)
            assert candidate is not None
            metrics = compare_to_reference(candidate, args.reference)
        print(
            "reference_metrics "
            + " ".join(f"{name}={value:.6f}" for name, value in metrics.items())
        )

    print(f"generated teeth={args.min_teeth}..{args.max_teeth} output={args.output}")


if __name__ == "__main__":
    main()