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_ascii_ply(path)
    if suffix in {".xyz", ".txt", ".csv"}:
        return np.loadtxt(path, dtype=np.float64)
    raise ValueError(f"Unsupported point cloud format: {path.suffix}")


def _read_ascii_ply(path: Path) -> np.ndarray:
    header_lines = 0
    vertex_count = None
    with path.open("r", encoding="ascii", errors="strict") as fh:
        for line in fh:
            header_lines += 1
            stripped = line.strip()
            if stripped.startswith("element vertex"):
                vertex_count = int(stripped.split()[-1])
            if stripped == "end_header":
                break
    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 vertex_count is not None and len(points) != vertex_count:
        raise ValueError(
            f"PLY vertex count mismatch: header={vertex_count}, data={len(points)}"
        )
    return points


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")