"""
I/O helpers for images and point-cloud files.
"""

from __future__ import annotations

from pathlib import Path
from typing import Optional, Tuple, Union

import numpy as np
from PIL import Image


def load_image(path: Union[str, Path, Image.Image, np.ndarray]) -> Image.Image:
    """
    Normalise any image input to a PIL RGB image.

    Parameters
    ----------
    path : str, Path, PIL.Image, or np.ndarray
        If a string/Path, loaded from disk.  If an ndarray, converted.

    Returns
    -------
    PIL.Image.Image
        RGB image ready for DepthPro.
    """
    if isinstance(path, (str, Path)):
        return Image.open(str(path)).convert("RGB")
    if isinstance(path, np.ndarray):
        if path.dtype != np.uint8:
            path = (path * 255).clip(0, 255).astype(np.uint8)
        return Image.fromarray(path).convert("RGB")
    if isinstance(path, Image.Image):
        return path.convert("RGB")
    raise TypeError(f"Unsupported image type: {type(path)}")


def save_point_cloud(
    path: Union[str, Path],
    points: np.ndarray,
    colors: Optional[np.ndarray] = None,
    normals: Optional[np.ndarray] = None,
) -> None:
    """
    Save a point cloud to a PLY file.

    Parameters
    ----------
    path : str or Path
        Output file path.  Should end with ``.ply``.
    points : np.ndarray
        (N, 3) float array of 3D positions.
    colors : np.ndarray, optional
        (N, 3) uint8 RGB colours.
    normals : np.ndarray, optional
        (N, 3) float array of normals.
    """
    path = Path(path)
    points = np.asarray(points, dtype=np.float32)
    has_colors = colors is not None
    has_normals = normals is not None

    header_lines = [
        "ply",
        "format ascii 1.0",
        f"element vertex {len(points)}",
        "property float x",
        "property float y",
        "property float z",
    ]

    if has_normals:
        header_lines += [
            "property float nx",
            "property float ny",
            "property float nz",
        ]

    if has_colors:
        header_lines += [
            "property uchar red",
            "property uchar green",
            "property uchar blue",
        ]

    header_lines += ["end_header"]

    with open(path, "w") as f:
        f.write("\n".join(header_lines) + "\n")
        for i in range(len(points)):
            row = [f"{points[i, 0]:.6f}", f"{points[i, 1]:.6f}", f"{points[i, 2]:.6f}"]
            if has_normals:
                row += [
                    f"{normals[i, 0]:.6f}",
                    f"{normals[i, 1]:.6f}",
                    f"{normals[i, 2]:.6f}",
                ]
            if has_colors:
                c = np.clip(colors[i], 0, 255).astype(np.uint8)
                row += [str(c[0]), str(c[1]), str(c[2])]
            f.write(" ".join(row) + "\n")