from __future__ import annotations

import csv
import json
import struct
from collections import Counter
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
LIVE_ROOT = ROOT / "preview_exports" / "live_update"
PARTS_DIR = LIVE_ROOT / "all_parts"
ASSEMBLY_DIR = LIVE_ROOT / "assembly"
VIEWER_ASSEMBLY_DIR = ROOT / "preview_exports" / "assembly"


COLORS = {
    "01 Truss": "#6f8798",
    "02 Floor": "#8ea85c",
    "03 Screen": "#a4adb7",
    "04 Building": "#c8b775",
    "05 Roof_Mounting": "#8c7567",
    "08 Tyipcal_Parts": "#5f7f9a",
}


def parse_stl(path: Path) -> dict:
    data = path.read_bytes()
    tri_count = struct.unpack_from("<I", data, 80)[0]
    expected = 84 + tri_count * 50
    if expected != len(data):
        raise RuntimeError(f"Unexpected STL size for {path.name}: expected {expected}, got {len(data)}")

    vertices: list[float] = []
    indices: list[int] = []
    vertex_map: dict[tuple[float, float, float], int] = {}
    offset = 84
    for _ in range(tri_count):
        offset += 12
        tri: list[int] = []
        for _ in range(3):
            vertex = struct.unpack_from("<3f", data, offset)
            offset += 12
            key = tuple(round(value, 6) for value in vertex)
            if key not in vertex_map:
                vertex_map[key] = len(vertices) // 3
                vertices.extend([float(key[0]), float(key[1]), float(key[2])])
            tri.append(vertex_map[key])
        indices.extend(tri)
        offset += 2

    xs = vertices[0::3] or [0]
    ys = vertices[1::3] or [0]
    zs = vertices[2::3] or [0]
    bounds = {
        "min": [min(xs), min(ys), min(zs)],
        "max": [max(xs), max(ys), max(zs)],
    }
    bounds["size"] = [bounds["max"][i] - bounds["min"][i] for i in range(3)]
    return {
        "format": "simple-indexed-triangle-mesh",
        "source": "Live Inventor update: IPT exported to STL, then parsed to JSON",
        "units": "millimeters",
        "triangleCount": tri_count,
        "vertexCount": len(vertices) // 3,
        "bounds": bounds,
        "vertices": vertices,
        "indices": indices,
    }


def transform_point(matrix: list[float], point: list[float]) -> list[float]:
    x, y, z = point
    return [
        matrix[0] * x + matrix[1] * y + matrix[2] * z + matrix[3],
        matrix[4] * x + matrix[5] * y + matrix[6] * z + matrix[7],
        matrix[8] * x + matrix[9] * y + matrix[10] * z + matrix[11],
    ]


def corners(bounds: dict) -> list[list[float]]:
    mn = bounds["min"]
    mx = bounds["max"]
    return [[x, y, z] for x in (mn[0], mx[0]) for y in (mn[1], mx[1]) for z in (mn[2], mx[2])]


def color_for(path: str) -> str:
    for prefix, color in COLORS.items():
        if path.startswith(prefix):
            return color
    return "#9da7b0"


def main() -> None:
    ASSEMBLY_DIR.mkdir(parents=True, exist_ok=True)
    part_rows = list(csv.DictReader((PARTS_DIR / "part-export-results.csv").open(newline="", encoding="utf-8-sig")))
    parts: dict[str, dict] = {}

    for row in part_rows:
        if row["status"] != "ok":
            continue
        try:
            mesh = parse_stl(PARTS_DIR / row["stl"])
        except Exception as exc:
            row["status"] = "error"
            row["error"] = f"STL conversion failed: {exc}"
            continue
        mesh.update({"id": row["id"], "part": Path(row["sourcePath"]).name, "sourcePath": row["sourcePath"]})
        mesh_name = row["id"] + ".mesh.json"
        (PARTS_DIR / mesh_name).write_text(json.dumps(mesh, separators=(",", ":")), encoding="utf-8")
        folder = row["sourcePath"].split("\\")[0]
        parts[row["id"]] = {
            "id": row["id"],
            "name": Path(row["sourcePath"]).name,
            "sourcePath": row["sourcePath"],
            "folder": folder,
            "meshJson": "../live_update/all_parts/" + mesh_name,
            "triangleCount": mesh["triangleCount"],
            "vertexCount": mesh["vertexCount"],
            "bounds": mesh["bounds"],
            "count": 0,
        }

    raw_occurrences = json.loads((ASSEMBLY_DIR / "leaf-occurrences.raw.json").read_text(encoding="utf-8-sig"))
    if isinstance(raw_occurrences, dict):
        raw_occurrences = [raw_occurrences]

    mins = [float("inf")] * 3
    maxs = [float("-inf")] * 3
    instances: list[dict] = []
    missing: set[str] = set()

    for index, occurrence in enumerate(raw_occurrences):
        part_id = occurrence["partId"]
        part = parts.get(part_id)
        if not part:
            missing.add(part_id)
            continue
        part["count"] += 1
        matrix = [float(value) for value in occurrence["transform"]]
        for corner in corners(part["bounds"]):
            point = transform_point(matrix, corner)
            for axis in range(3):
                mins[axis] = min(mins[axis], point[axis])
                maxs[axis] = max(maxs[axis], point[axis])
        source_path = occurrence["sourcePath"]
        instances.append(
            {
                "id": f"occ_{index:04d}",
                "name": occurrence["occurrenceName"],
                "partId": part_id,
                "sourcePath": source_path,
                "folder": source_path.split("\\")[0],
                "matrixRowMajor": matrix,
                "color": color_for(source_path),
            }
        )

    bounds = {"min": mins, "max": maxs, "size": [maxs[i] - mins[i] for i in range(3)]}
    folder_counts = Counter(instance["folder"] for instance in instances)
    assembly = {
        "name": "ROOF_MOUNT_PLATFORM_LIVE",
        "sourceAssembly": "ROOF_MOUNT_PLATFORM.iam",
        "units": "millimeters",
        "basis": "Live rebuild from submitted parameters using Inventor API, iLogic RunRules attempt, Update2, STL export, and JSON conversion.",
        "partCount": len(parts),
        "instanceCount": len(instances),
        "missingMeshPartIds": sorted(missing),
        "estimatedInstancedTriangles": sum(parts[i["partId"]]["triangleCount"] for i in instances),
        "bounds": bounds,
        "folderInstanceCounts": dict(sorted(folder_counts.items())),
        "parts": parts,
        "instances": instances,
    }

    live_catalog = VIEWER_ASSEMBLY_DIR / "roof_mount_platform_live_assembly.json"
    live_catalog.write_text(json.dumps(assembly, separators=(",", ":")), encoding="utf-8")
    (ASSEMBLY_DIR / "roof_mount_platform_live_assembly.json").write_text(json.dumps(assembly, indent=2), encoding="utf-8")

    with (ASSEMBLY_DIR / "live-assembly-summary.csv").open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["count", "folder", "partId", "sourcePath", "triangles", "vertices"])
        for part in sorted(parts.values(), key=lambda item: (item["folder"], -item["count"], item["sourcePath"].lower())):
            writer.writerow([part["count"], part["folder"], part["id"], part["sourcePath"], part["triangleCount"], part["vertexCount"]])

    print(json.dumps({"ok": True, "parts": len(parts), "instances": len(instances), "catalog": str(live_catalog)}, indent=2))


if __name__ == "__main__":
    main()