from __future__ import annotations

import argparse
import html
from pathlib import Path

import cv2
import numpy as np


IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp", ".webp"}


def read_gray(path: Path) -> np.ndarray:
    raw = np.fromfile(str(path), dtype=np.uint8)
    image = cv2.imdecode(raw, cv2.IMREAD_GRAYSCALE)
    if image is None:
        raise ValueError(f"Cannot read image: {path}")
    return image


def write_image(path: Path, image: np.ndarray) -> None:
    ok, encoded = cv2.imencode(path.suffix, image)
    if not ok:
        raise ValueError(f"Cannot encode image: {path}")
    encoded.tofile(str(path))


def display_to_mask(gray: np.ndarray) -> np.ndarray:
    return (gray < 128).astype(np.uint8) * 255


def mask_to_display(mask: np.ndarray) -> np.ndarray:
    return np.where(mask > 0, 0, 255).astype(np.uint8)


def foreground_bbox(mask: np.ndarray) -> tuple[int, int, int, int]:
    ys, xs = np.where(mask > 0)
    if xs.size == 0:
        height, width = mask.shape
        return 0, 0, width, height
    return int(xs.min()), int(ys.min()), int(xs.max() + 1), int(ys.max() + 1)


def contour_complexity(mask: np.ndarray) -> float:
    x1, y1, x2, y2 = foreground_bbox(mask)
    bbox_area = max(1, (x2 - x1) * (y2 - y1))
    contours, _ = cv2.findContours(mask, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
    perimeter = sum(cv2.arcLength(contour, True) for contour in contours)
    return float(perimeter / np.sqrt(bbox_area))


def vector_params(mask: np.ndarray) -> tuple[int, float]:
    complexity = contour_complexity(mask)
    if complexity >= 28.0:
        return 1, 0.0009
    if complexity >= 18.0:
        return 1, 0.0014
    return 2, 0.0022


def smooth_mask(mask: np.ndarray, iterations: int) -> np.ndarray:
    smoothed = mask.copy()
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    for _ in range(iterations):
        blurred = cv2.GaussianBlur(smoothed, (3, 3), 0)
        smoothed = np.where(blurred >= 128, 255, 0).astype(np.uint8)
        smoothed = cv2.morphologyEx(smoothed, cv2.MORPH_CLOSE, kernel)
    return smoothed


def points_to_smooth_path(points: np.ndarray) -> str:
    if len(points) < 3:
        return ""

    midpoints = (points + np.roll(points, -1, axis=0)) * 0.5
    commands = [f"M {midpoints[-1, 0]:.2f} {midpoints[-1, 1]:.2f}"]
    for point, midpoint in zip(points, midpoints):
        commands.append(
            f"Q {point[0]:.2f} {point[1]:.2f} {midpoint[0]:.2f} {midpoint[1]:.2f}"
        )
    commands.append("Z")
    return " ".join(commands)


def contour_to_path(contour: np.ndarray, epsilon_ratio: float) -> str:
    perimeter = cv2.arcLength(contour, True)
    epsilon = max(0.35, perimeter * epsilon_ratio)
    approx = cv2.approxPolyDP(contour, epsilon, True).reshape(-1, 2)
    if len(approx) < 3:
        return ""

    points = approx.astype(np.float64)
    return points_to_smooth_path(points)


def mask_to_svg(mask: np.ndarray, title: str, epsilon_ratio: float) -> str:
    height, width = mask.shape
    contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    contours = sorted(contours, key=cv2.contourArea, reverse=True)
    x1, y1, x2, y2 = foreground_bbox(mask)
    bbox_area = max(1, (x2 - x1) * (y2 - y1))
    min_area = max(2.0, bbox_area * 0.000008)

    paths = []
    for contour in contours:
        if abs(cv2.contourArea(contour)) < min_area:
            continue
        path = contour_to_path(contour, epsilon_ratio=epsilon_ratio)
        if path:
            paths.append(path)

    escaped_title = html.escape(title)
    path_data = " ".join(paths)
    return "\n".join(
        [
            '<?xml version="1.0" encoding="UTF-8"?>',
            f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" '
            f'viewBox="0 0 {width} {height}" role="img" aria-labelledby="title">',
            f"  <title>{escaped_title}</title>",
            '  <rect width="100%" height="100%" fill="white"/>',
            f'  <path d="{path_data}" fill="black" fill-rule="evenodd"/>',
            "</svg>",
            "",
        ]
    )


def iter_images(input_dir: Path) -> list[Path]:
    return sorted(
        path
        for path in input_dir.iterdir()
        if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES
    )


def process_image(
    image_path: Path,
    smoothed_dir: Path,
    svg_dir: Path,
    smooth_iterations: int | None,
    epsilon_ratio: float | None,
) -> None:
    gray = read_gray(image_path)
    mask = display_to_mask(gray)
    auto_smooth_iterations, auto_epsilon_ratio = vector_params(mask)
    iterations = (
        auto_smooth_iterations if smooth_iterations is None else smooth_iterations
    )
    ratio = auto_epsilon_ratio if epsilon_ratio is None else epsilon_ratio

    smoothed = smooth_mask(mask, iterations=iterations)
    stem = image_path.stem.removesuffix("_cleaned")
    write_image(smoothed_dir / f"{stem}_smoothed.png", mask_to_display(smoothed))
    svg = mask_to_svg(smoothed, title=stem, epsilon_ratio=ratio)
    (svg_dir / f"{stem}.svg").write_text(svg, encoding="utf-8")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Smooth and vectorize cleaned complex seal-script characters."
    )
    parser.add_argument(
        "--input-dir",
        type=Path,
        default=Path("PreProcessing_Common/cleaned_characters"),
    )
    parser.add_argument(
        "--smoothed-dir",
        type=Path,
        default=Path("PreProcessing_Common/smoothed_characters"),
    )
    parser.add_argument(
        "--svg-dir",
        type=Path,
        default=Path("PreProcessing_Common/svg_characters"),
    )
    parser.add_argument(
        "--smooth-iterations",
        type=int,
        default=None,
        help="Override automatic smoothing. Dense characters default to less smoothing.",
    )
    parser.add_argument(
        "--epsilon-ratio",
        type=float,
        default=None,
        help="Override contour simplification ratio. Smaller values keep more detail.",
    )
    return parser


def main() -> None:
    args = build_parser().parse_args()
    args.smoothed_dir.mkdir(parents=True, exist_ok=True)
    args.svg_dir.mkdir(parents=True, exist_ok=True)
    image_paths = iter_images(args.input_dir)
    if not image_paths:
        raise SystemExit(f"No images found in {args.input_dir}")

    for image_path in image_paths:
        process_image(
            image_path=image_path,
            smoothed_dir=args.smoothed_dir,
            svg_dir=args.svg_dir,
            smooth_iterations=args.smooth_iterations,
            epsilon_ratio=args.epsilon_ratio,
        )
        print(f"vectorized: {image_path.name}")


if __name__ == "__main__":
    main()