from __future__ import annotations

import csv
import json
from pathlib import Path
from typing import Iterable

import numpy as np


def read_point_cloud(path: str | Path) -> np.ndarray:
    path = Path(path)
    suffix = path.suffix.lower()
    if suffix == ".ply":
        return _read_ply(path)
    if suffix in {".xyz", ".txt", ".csv"}:
        return np.loadtxt(path, dtype=np.float64)
    raise ValueError(f"Unsupported point cloud format: {path.suffix}")


# PLY scalar property type -> numpy dtype code.
_PLY_TYPE_TO_NP = {
    "char": "i1",
    "int8": "i1",
    "uchar": "u1",
    "uint8": "u1",
    "short": "i2",
    "int16": "i2",
    "ushort": "u2",
    "uint16": "u2",
    "int": "i4",
    "int32": "i4",
    "uint": "u4",
    "uint32": "u4",
    "float": "f4",
    "float32": "f4",
    "double": "f8",
    "float64": "f8",
}


def _parse_ply_header(
    path: Path,
) -> tuple[str, int, int, int, list[tuple[str, str]]]:
    """Parse a PLY header.

    Returns ``(format, header_lines, header_bytes, vertex_count,
    vertex_props)`` where ``format`` is one of ``ascii`` /
    ``binary_little_endian`` / ``binary_big_endian`` and ``vertex_props`` is
    the ordered ``(type, name)`` list for the vertex element. The file is
    read in binary so the byte offset of the data block is exact (needed for
    binary PLY).
    """
    header_lines = 0
    header_bytes = 0
    vertex_count: int | None = None
    vertex_props: list[tuple[str, str]] = []
    in_vertex = False
    fmt = "ascii"
    with path.open("rb") as fh:
        first = fh.readline()
        header_bytes += len(first)
        header_lines += 1
        if first.strip().lower() != b"ply":
            raise ValueError(f"Not a PLY file (missing 'ply' magic): {path}")
        while True:
            raw = fh.readline()
            if not raw:
                raise ValueError(f"PLY header missing end_header: {path}")
            header_bytes += len(raw)
            header_lines += 1
            stripped = raw.decode("ascii", errors="replace").strip()
            if not stripped:
                continue
            tokens = stripped.split()
            keyword = tokens[0].lower()
            if keyword == "format":
                fmt = tokens[1].lower()
            elif keyword == "element":
                in_vertex = tokens[1].lower() == "vertex"
                if in_vertex:
                    vertex_count = int(tokens[2])
            elif keyword == "property" and in_vertex:
                if tokens[1].lower() == "list":
                    vertex_props.append(("list", tokens[-1]))
                else:
                    vertex_props.append((tokens[1].lower(), tokens[-1]))
            elif keyword == "end_header":
                break
    if vertex_count is None:
        raise ValueError(f"PLY file has no vertex element: {path}")
    return fmt, header_lines, header_bytes, vertex_count, vertex_props


def _read_ply(path: Path) -> np.ndarray:
    """Read a PLY point cloud (ASCII or binary) into an ``(N, 3)`` array."""
    fmt, header_lines, header_bytes, vertex_count, vertex_props = _parse_ply_header(path)
    if fmt == "ascii":
        points = np.loadtxt(path, skiprows=header_lines, dtype=np.float64)
        if points.ndim == 1:
            points = points[None, :]
        if points.shape[1] < 3:
            raise ValueError(f"PLY file does not contain XYZ coordinates: {path}")
        points = points[:, :3]
        if len(points) != vertex_count:
            raise ValueError(
                f"PLY vertex count mismatch: header={vertex_count}, data={len(points)}"
            )
        return points
    if fmt in {"binary_little_endian", "binary_big_endian"}:
        return _read_binary_ply(path, header_bytes, vertex_count, vertex_props, fmt)
    raise ValueError(f"Unsupported PLY format '{fmt}': {path}")


def _read_binary_ply(
    path: Path,
    header_bytes: int,
    vertex_count: int,
    vertex_props: list[tuple[str, str]],
    fmt: str,
) -> np.ndarray:
    if any(ptype == "list" for ptype, _ in vertex_props):
        raise ValueError(
            f"Binary PLY with list properties in the vertex element is unsupported: {path}"
        )
    missing = {"x", "y", "z"} - {name for _, name in vertex_props}
    if missing:
        raise ValueError(
            f"PLY vertex element missing coordinate field(s) {sorted(missing)}: {path}"
        )
    endian = "<" if fmt == "binary_little_endian" else ">"
    try:
        record_dtype = np.dtype(
            [(name, endian + _PLY_TYPE_TO_NP[ptype]) for ptype, name in vertex_props]
        )
    except KeyError as exc:
        raise ValueError(f"Unsupported PLY property type {exc} in {path}") from exc
    expected = record_dtype.itemsize * vertex_count
    with path.open("rb") as fh:
        fh.seek(header_bytes)
        buffer = fh.read(expected)
    if len(buffer) < expected:
        raise ValueError(
            f"Binary PLY truncated: expected {expected} bytes of vertex data, "
            f"got {len(buffer)}: {path}"
        )
    records = np.frombuffer(buffer, dtype=record_dtype, count=vertex_count)
    return np.column_stack([records["x"], records["y"], records["z"]]).astype(np.float64)


# Backwards-compatible alias. ``cad_io.py`` imports this name; it now reads
# both ASCII and binary PLY transparently.
_read_ascii_ply = _read_ply


def read_point_cloud_cropped(
    path: str | Path,
    aabb: dict,
    *,
    chunk_vertices: int = 4_000_000,
) -> np.ndarray:
    """Stream a (possibly huge) PLY and return only points inside ``aabb``.

    For multi-GB clouds (hundreds of millions / billions of points) loading
    the whole file is infeasible. This reads the vertex block in chunks,
    keeps only points within the axis-aligned bbox, and concatenates the
    survivors. Falls back to a full read + mask for ASCII / text formats.

    ``aabb`` keys: ``x_min/x_max/y_min/y_max/z_min/z_max`` (same frame as the
    file). Returns an ``(M, 3) float64`` array.
    """
    path = Path(path)
    suffix = path.suffix.lower()
    bounds = (
        float(aabb["x_min"]), float(aabb["x_max"]),
        float(aabb["y_min"]), float(aabb["y_max"]),
        float(aabb["z_min"]), float(aabb["z_max"]),
    )
    if suffix != ".ply":
        pts = read_point_cloud(path)
        return _crop_points(pts, bounds)

    fmt, _, header_bytes, vertex_count, vertex_props = _parse_ply_header(path)
    if fmt == "ascii":
        pts = _read_ply(path)
        return _crop_points(pts, bounds)
    if fmt not in {"binary_little_endian", "binary_big_endian"}:
        raise ValueError(f"Unsupported PLY format '{fmt}': {path}")
    if any(ptype == "list" for ptype, _ in vertex_props):
        raise ValueError(f"Binary PLY with list vertex properties unsupported: {path}")
    missing = {"x", "y", "z"} - {name for _, name in vertex_props}
    if missing:
        raise ValueError(f"PLY missing coordinate field(s) {sorted(missing)}: {path}")

    endian = "<" if fmt == "binary_little_endian" else ">"
    record_dtype = np.dtype(
        [(name, endian + _PLY_TYPE_TO_NP[ptype]) for ptype, name in vertex_props]
    )
    record_bytes = record_dtype.itemsize
    kept_chunks: list[np.ndarray] = []
    with path.open("rb") as fh:
        fh.seek(header_bytes)
        remaining = vertex_count
        while remaining > 0:
            n = min(chunk_vertices, remaining)
            buf = fh.read(n * record_bytes)
            if len(buf) < n * record_bytes:
                n = len(buf) // record_bytes
                if n == 0:
                    break
                buf = buf[: n * record_bytes]
            recs = np.frombuffer(buf, dtype=record_dtype, count=n)
            xyz = np.column_stack(
                [recs["x"], recs["y"], recs["z"]]
            ).astype(np.float64)
            crop = _crop_points(xyz, bounds)
            if len(crop):
                kept_chunks.append(crop)
            remaining -= n
    if not kept_chunks:
        return np.zeros((0, 3), dtype=np.float64)
    return np.concatenate(kept_chunks, axis=0)


def _crop_points(points: np.ndarray, bounds: tuple) -> np.ndarray:
    x_lo, x_hi, y_lo, y_hi, z_lo, z_hi = bounds
    mask = (
        (points[:, 0] >= x_lo) & (points[:, 0] <= x_hi)
        & (points[:, 1] >= y_lo) & (points[:, 1] <= y_hi)
        & (points[:, 2] >= z_lo) & (points[:, 2] <= z_hi)
    )
    return points[mask]


def ensure_output_dir(path: str | Path) -> Path:
    out_dir = Path(path)
    out_dir.mkdir(parents=True, exist_ok=True)
    return out_dir


def write_json(path: str | Path, payload: dict) -> None:
    Path(path).write_text(
        json.dumps(payload, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )


def write_csv(path: str | Path, fieldnames: list[str], rows: Iterable[dict]) -> None:
    with Path(path).open("w", encoding="utf-8", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=fieldnames)
        writer.writeheader()
        for row in rows:
            writer.writerow(row)


def write_colored_ply(
    path: str | Path,
    points: np.ndarray,
    colors: np.ndarray,
) -> None:
    path = Path(path)
    if len(points) != len(colors):
        raise ValueError("Point/color length mismatch for PLY export.")
    points = np.asarray(points, dtype=np.float64)
    colors = np.asarray(colors, dtype=np.uint8)
    with path.open("w", encoding="ascii", newline="\n") as fh:
        fh.write("ply\n")
        fh.write("format ascii 1.0\n")
        fh.write(f"element vertex {len(points)}\n")
        fh.write("property float x\n")
        fh.write("property float y\n")
        fh.write("property float z\n")
        fh.write("property uchar red\n")
        fh.write("property uchar green\n")
        fh.write("property uchar blue\n")
        fh.write("end_header\n")
        for (x, y, z), (r, g, b) in zip(points, colors, strict=True):
            fh.write(f"{x:.6f} {y:.6f} {z:.6f} {int(r)} {int(g)} {int(b)}\n")