"""Build and preview the AC28 supervised CAD-to-scan dataset.

Each manifest record uses the following contract:

* ``inputs``: the shared CATPart design file and one production Excel file;
* ``target``: the corresponding merged point-cloud PLY and its six scan views;
* ``label``: defect coordinates from the CSV, or ``null`` for an OK unit.

The source directory is deliberately kept external to the Pixi project so the
manifest contains portable paths relative to the data root.
"""

from __future__ import annotations

import argparse
import csv
import json
import re
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.patches import FancyBboxPatch


CATIA_DIR = "CATIA File (설계 데이터)"
EXCEL_DIR = "Excel File (성형정보)"
CSV_DIR = "CSV File (불량 좌표값)"
PLY_DIR = "PLY File (측정 데이터)"
CATPART_NAME = "AC28_RR DT UPR LHCUR.CATPart"
PRODUCTION_RE = re.compile(
    r"_(?P<date>\d{8}) (?P<time>\d+시(?:\d+분)?) 생산_(?P<status>NG|OK)$"
)


@dataclass(frozen=True)
class ProductionKey:
    date: str
    time: str
    status: str

    @property
    def sample_id(self) -> str:
        return f"AC28_{self.date}_{self.time}_{self.status}"


def production_key(name: str) -> ProductionKey:
    """Extract the timestamp/status portion shared across Excel, PLY, and CSV."""
    match = PRODUCTION_RE.search(Path(name).stem)
    if not match:
        raise ValueError(f"Could not parse production date/status from {name!r}")
    return ProductionKey(**match.groupdict())


def relative(path: Path, root: Path) -> str:
    return path.relative_to(root).as_posix()


def display_time(value: str) -> str:
    """Convert Korean filename time text to an ASCII-safe plot label."""
    match = re.fullmatch(r"(\d+)시(?:(\d+)분)?", value)
    if not match:
        return value
    return f"{int(match.group(1)):02d}:{int(match.group(2) or 0):02d}"


def read_ply_header(path: Path) -> dict[str, Any]:
    """Read only the binary PLY header, never its multi-million-point body."""
    with path.open("rb") as handle:
        header_lines: list[bytes] = []
        while True:
            line = handle.readline()
            if not line:
                raise ValueError(f"PLY header is incomplete: {path}")
            header_lines.append(line)
            if line.strip() == b"end_header":
                break
        byte_offset = handle.tell()

    header = b"".join(header_lines).decode("ascii", errors="strict")
    vertex_match = re.search(r"^element vertex (\d+)$", header, flags=re.MULTILINE)
    if not vertex_match:
        raise ValueError(f"PLY has no vertex count: {path}")
    if "format binary_little_endian 1.0" not in header:
        raise ValueError(f"Expected binary little-endian PLY: {path}")

    return {
        "format": "binary_little_endian",
        "vertex_count": int(vertex_match.group(1)),
        "header_bytes": byte_offset,
    }


def read_excel_metadata(path: Path) -> dict[str, Any]:
    """Capture a small, JSON-safe preview without putting sheet contents in labels."""
    workbook = pd.ExcelFile(path, engine="xlrd")
    first_sheet = workbook.sheet_names[0]
    preview = pd.read_excel(path, sheet_name=first_sheet, header=None, nrows=4, engine="xlrd")
    values = preview.where(pd.notna(preview), None).values.tolist()
    return {"sheets": workbook.sheet_names, "preview_rows": values}


def read_label(path: Path) -> list[dict[str, Any]]:
    with path.open("r", encoding="utf-8-sig", newline="") as handle:
        rows = list(csv.DictReader(handle))
    labels: list[dict[str, Any]] = []
    for row in rows:
        labels.append(
            {
                "id": int(row["No"]),
                "x": int(row["X"]),
                "y": int(row["Y"]),
                "width": int(row["W"]) if row.get("W") else None,
                "height": int(row["H"]) if row.get("H") else None,
                "error": row.get("Error") or None,
            }
        )
    return labels


def csv_by_key(root: Path) -> dict[ProductionKey, Path]:
    result: dict[ProductionKey, Path] = {}
    for path in sorted((root / CSV_DIR).glob("*.csv")):
        key = production_key(path.name)
        if key in result:
            raise ValueError(f"Duplicate CSV label for {key}: {path} and {result[key]}")
        result[key] = path
    return result


def build_records(root: Path) -> list[dict[str, Any]]:
    catpart = root / CATIA_DIR / CATPART_NAME
    if not catpart.is_file():
        raise FileNotFoundError(catpart)

    labels = csv_by_key(root)
    records: list[dict[str, Any]] = []
    for excel_path in sorted((root / EXCEL_DIR).glob("*.Xls")):
        key = production_key(excel_path.name)
        ply_folder = root / PLY_DIR / excel_path.stem
        if not ply_folder.is_dir():
            raise FileNotFoundError(f"No PLY folder matching {excel_path.name}: {ply_folder}")

        ply_paths = sorted(ply_folder.glob("*.ply"))
        primary_ply = ply_folder / "all_merge_pointcloud_full.ply"
        if not primary_ply.is_file():
            raise FileNotFoundError(f"Missing primary merged PLY: {primary_ply}")

        label_path = labels.get(key)
        if key.status == "NG" and label_path is None:
            raise ValueError(f"NG record has no defect-coordinate CSV: {excel_path.name}")
        if key.status == "OK" and label_path is not None:
            raise ValueError(f"OK record unexpectedly has a defect CSV: {excel_path.name}")

        records.append(
            {
                "sample_id": key.sample_id,
                "production": asdict(key),
                "inputs": {
                    "catpart": relative(catpart, root),
                    "excel": relative(excel_path, root),
                    "excel_metadata": read_excel_metadata(excel_path),
                },
                "target": {
                    "ply_directory": relative(ply_folder, root),
                    "primary_ply": relative(primary_ply, root),
                    "primary_ply_header": read_ply_header(primary_ply),
                    "scan_plys": [
                        {"path": relative(path, root), "header": read_ply_header(path)}
                        for path in ply_paths
                    ],
                },
                "label": (
                    {
                        "csv": relative(label_path, root),
                        "defects": read_label(label_path),
                        "coordinate_system": "source inspection-image pixels; registration to PLY is not supplied",
                    }
                    if label_path
                    else None
                ),
            }
        )
    return records


def write_manifest(records: list[dict[str, Any]], output_dir: Path) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    manifest = output_dir / "ac28_matching_dataset.jsonl"
    with manifest.open("w", encoding="utf-8") as handle:
        for record in records:
            handle.write(json.dumps(record, ensure_ascii=False) + "\n")

    summary = {
        "schema": {"inputs": ["CATPart", "Excel"], "target": "PLY", "label": "CSV"},
        "record_count": len(records),
        "ng_count": sum(record["production"]["status"] == "NG" for record in records),
        "ok_count": sum(record["production"]["status"] == "OK" for record in records),
        "shared_catpart": records[0]["inputs"]["catpart"],
        "records": [
            {
                "sample_id": record["sample_id"],
                "excel": record["inputs"]["excel"],
                "primary_ply": record["target"]["primary_ply"],
                "label_csv": record["label"]["csv"] if record["label"] else None,
            }
            for record in records
        ],
    }
    (output_dir / "ac28_matching_dataset_summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )


def sample_point_cloud(path: Path, header: dict[str, Any], maximum_points: int = 8000) -> np.ndarray:
    """Uniformly sample RGB points from a large binary PLY through a memory map."""
    dtype = np.dtype(
        [("x", "<f4"), ("y", "<f4"), ("z", "<f4"), ("r", "u1"), ("g", "u1"), ("b", "u1")]
    )
    count = header["vertex_count"]
    indices = np.linspace(0, count - 1, min(count, maximum_points), dtype=np.int64)
    points = np.memmap(path, dtype=dtype, mode="r", offset=header["header_bytes"], shape=(count,))
    sampled = points[indices].copy()
    del points
    return sampled


def render_preview(records: list[dict[str, Any]], root: Path, output_dir: Path) -> Path:
    output_dir.mkdir(parents=True, exist_ok=True)
    figure, axes = plt.subplots(2, 2, figsize=(14, 12), constrained_layout=True)
    for axis, record in zip(axes.flat, records, strict=True):
        header = record["target"]["primary_ply_header"]
        points = sample_point_cloud(root / record["target"]["primary_ply"], header)
        colors = np.column_stack((points["r"], points["g"], points["b"])) / 255.0
        axis.scatter(points["x"], points["y"], c=colors, s=0.3, linewidths=0, rasterized=True)
        axis.set_aspect("equal", adjustable="box")
        axis.set_xlabel("PLY x")
        axis.set_ylabel("PLY y")
        status = record["production"]["status"]
        production = record["production"]
        title = (
            f"AC28 {production['date']} {display_time(production['time'])} {status}\n"
            f"{header['vertex_count']:,} vertices (8,000 sampled)"
        )
        if record["label"]:
            defect = record["label"]["defects"][0]
            title += f"\nCSV defect #{defect['id']} @ image ({defect['x']}, {defect['y']})"
            axis.set_title(title, color="crimson", fontsize=10)
        else:
            title += "\nCSV label: none (OK production)"
            axis.set_title(title, color="seagreen", fontsize=10)
        axis.grid(alpha=0.2)

    preview = output_dir / "ac28_pointcloud_label_preview.png"
    figure.suptitle("AC28: CATPart + Excel → PLY target; CSV defect labels", fontsize=15)
    figure.savefig(preview, dpi=180, bbox_inches="tight")
    plt.close(figure)
    return preview


def record_title(record: dict[str, Any]) -> str:
    production = record["production"]
    return f"AC28 {production['date']} {display_time(production['time'])} {production['status']}"


def render_profile_preview(records: list[dict[str, Any]], root: Path, output_dir: Path) -> Path:
    """Render a complementary x-z projection; colour represents elevation."""
    output_dir.mkdir(parents=True, exist_ok=True)
    figure, axes = plt.subplots(2, 2, figsize=(14, 12), constrained_layout=True)
    for axis, record in zip(axes.flat, records, strict=True):
        header = record["target"]["primary_ply_header"]
        points = sample_point_cloud(root / record["target"]["primary_ply"], header)
        axis.scatter(
            points["x"],
            points["z"],
            c=points["z"],
            cmap="viridis",
            s=0.45,
            linewidths=0,
            rasterized=True,
        )
        axis.set_aspect("equal", adjustable="box")
        axis.set_xlabel("PLY x")
        axis.set_ylabel("PLY z / elevation")
        axis.set_title(record_title(record), color="crimson" if record["label"] else "seagreen", fontsize=11)
        axis.grid(alpha=0.2)

    preview = output_dir / "ac28_pointcloud_profile_preview.png"
    figure.suptitle("AC28 merged point-cloud profile views (uniform 8,000-point samples)", fontsize=15)
    figure.savefig(preview, dpi=180, bbox_inches="tight")
    plt.close(figure)
    return preview


def render_label_preview(records: list[dict[str, Any]], output_dir: Path) -> Path:
    """Show the CSV labels in their native inspection-image coordinate system."""
    output_dir.mkdir(parents=True, exist_ok=True)
    figure, axis = plt.subplots(figsize=(10, 8), constrained_layout=True)
    labels = [
        (record, defect)
        for record in records
        if record["label"]
        for defect in record["label"]["defects"]
    ]
    max_x = max(defect["x"] for _, defect in labels)
    max_y = max(defect["y"] for _, defect in labels)
    axis.set_xlim(0, max(2048, max_x + 300))
    axis.set_ylim(max(10240, max_y + 700), 0)  # image coordinates grow downward
    colours = ["#d62728", "#ff7f0e", "#9467bd"]
    offsets = [(12, -34), (12, -18), (12, 38)]
    for colour, offset, (record, defect) in zip(colours, offsets, labels, strict=True):
        axis.scatter(defect["x"], defect["y"], s=170, color=colour, edgecolors="black", zorder=3)
        axis.annotate(
            f"{record_title(record)}\ndefect #{defect['id']} ({defect['x']}, {defect['y']})",
            (defect["x"], defect["y"]),
            xytext=offset,
            textcoords="offset points",
            fontsize=9,
        )
    axis.set_title("NG defect labels in inspection-image coordinates")
    axis.set_xlabel("CSV X (pixels)")
    axis.set_ylabel("CSV Y (pixels; origin at upper-left)")
    axis.grid(alpha=0.3)
    preview = output_dir / "ac28_csv_label_coordinates.png"
    figure.savefig(preview, dpi=180, bbox_inches="tight")
    plt.close(figure)
    return preview


def render_contract_preview(output_dir: Path) -> Path:
    """Render the model-ready data contract as a simple data-flow diagram."""
    output_dir.mkdir(parents=True, exist_ok=True)
    figure, axis = plt.subplots(figsize=(13, 4), constrained_layout=True)
    axis.set_axis_off()
    boxes = [
        (0.03, 0.22, "Input X", "CATPart (shared)\n+ Excel (per unit)", "#dbeafe"),
        (0.30, 0.21, "Match key", "date + time +\nproduction status", "#fef3c7"),
        (0.57, 0.22, "Target Y", "merged PLY +\n6 scan-view PLYs", "#dcfce7"),
        (0.84, 0.14, "Label", "CSV defect pixels\n(NG only)", "#fee2e2"),
    ]
    for x, width, heading, body, colour in boxes:
        patch = FancyBboxPatch(
            (x, 0.28), width, 0.44, boxstyle="round,pad=0.02", facecolor=colour, edgecolor="#334155", linewidth=1.5
        )
        axis.add_patch(patch)
        axis.text(x + width / 2, 0.60, heading, ha="center", va="center", fontsize=12, fontweight="bold")
        axis.text(x + width / 2, 0.42, body, ha="center", va="center", fontsize=10)
    for start, end in ((0.25, 0.30), (0.52, 0.57), (0.80, 0.84)):
        axis.annotate("", xy=(end, 0.5), xytext=(start, 0.5), arrowprops={"arrowstyle": "->", "lw": 2})
    axis.text(0.5, 0.92, "AC28 supervised matching dataset", ha="center", va="center", fontsize=16, fontweight="bold")
    axis.text(
        0.5,
        0.08,
        "CSV coordinates are inspection-image pixels; a camera-to-PLY registration is required for vertex-level labels.",
        ha="center",
        va="center",
        fontsize=10,
        color="#475569",
    )
    preview = output_dir / "ac28_dataset_contract.png"
    figure.savefig(preview, dpi=180, bbox_inches="tight")
    plt.close(figure)
    return preview


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data-root", type=Path, default=Path("/workspace/moltex-sample"))
    parser.add_argument("--output-dir", type=Path, default=Path("artifacts"))
    parser.add_argument("--skip-preview", action="store_true")
    args = parser.parse_args()

    records = build_records(args.data_root)
    write_manifest(records, args.output_dir)
    if not args.skip_preview:
        previews = [
            render_preview(records, args.data_root, args.output_dir),
            render_profile_preview(records, args.data_root, args.output_dir),
            render_label_preview(records, args.output_dir),
            render_contract_preview(args.output_dir),
        ]
        for preview in previews:
            print(f"Wrote preview: {preview}")
    print(f"Wrote {len(records)} matched records to {args.output_dir / 'ac28_matching_dataset.jsonl'}")


if __name__ == "__main__":
    main()