# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 Lucasgitguy-co
import hashlib
import json
from pathlib import Path

import numpy as np
import torch
from PIL import Image, ImageOps
from torch.utils.data import Dataset

from .splats import Splats, read_ply


def grid_centers(grid):
    axis = (np.arange(grid, dtype=np.float32) + .5) / grid - .5
    return np.stack(np.meshgrid(axis, axis, axis, indexing="ij"), axis=-1).reshape(-1, 3)


def cell_indices(xyz, grid):
    if not np.isfinite(xyz).all() or (np.abs(xyz) > .50001).any():
        raise ValueError("Targets must be aligned inside [-0.5,0.5]^3. No silent normalization is performed.")
    ijk = np.floor((xyz + .5) * grid).astype(np.int64).clip(0, grid-1)
    return (ijk[:, 0] * grid + ijk[:, 1]) * grid + ijk[:, 2]


def splat_targets(splats, grid):
    """One highest-opacity teacher Gaussian per occupied cell; explicitly lossy."""
    cells = cell_indices(splats.xyz, grid)
    order = np.argsort(-splats.opacity[:, 0], kind="stable")
    _, first = np.unique(cells[order], return_index=True)
    keep = order[first]
    slots = cells[keep]
    result = {"occ": np.zeros((grid**3, 1), np.float32), "xyz": grid_centers(grid),
              "rgb": np.zeros((grid**3, 3), np.float32),
              "scale": np.full((grid**3, 3), 1 / grid / 3, np.float32),
              "quat": np.tile(np.array([1, 0, 0, 0], np.float32), (grid**3, 1)),
              "opacity": np.zeros((grid**3, 1), np.float32),
              "geometry": np.zeros((grid**3, 2), np.float32),
              "geometry_valid": np.zeros(2, np.float32), "grid": np.array(grid, np.int32)}
    result["occ"][slots] = 1
    for key in ("xyz", "rgb", "scale", "quat", "opacity"):
        result[key][slots] = getattr(splats, key)[keep]
    return result


def mesh_occupancy(path, transform_path, grid):
    import trimesh
    if not transform_path:
        raise ValueError("Every geometry teacher needs an explicit --*-transform JSON matrix, even for identity.")
    transform = np.asarray(json.loads(Path(transform_path).read_text()), dtype=float)
    if transform.shape != (4, 4) or not np.isfinite(transform).all() or not np.allclose(transform[3], [0,0,0,1]):
        raise ValueError("Transform must be a finite affine 4x4 matrix")
    mesh = trimesh.load(str(path), force="mesh", process=False)
    if not isinstance(mesh, trimesh.Trimesh) or len(mesh.faces) == 0:
        raise ValueError("Geometry teacher must contain a triangle mesh")
    mesh.apply_transform(transform)
    cell_indices(np.asarray(mesh.vertices), grid)  # validates the whole object
    points, _ = trimesh.sample.sample_surface(mesh, min(200_000, grid**3 * 16), seed=42)
    mask = np.zeros(grid**3, np.float32)
    mask[np.unique(cell_indices(points, grid))] = 1
    return mask


def cache_sample(root, sample_id, image_path, splats, grid=16, source="idol",
                 coarse=None, coarse_transform=None, fine=None, fine_transform=None, extra=None):
    root = Path(root)
    if not sample_id or any(c not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_" for c in sample_id):
        raise ValueError("Sample ID may contain only letters, digits, - and _")
    if grid not in (8, 16, 32):
        raise ValueError("Grid must be 8, 16 or 32")
    target = splat_targets(splats, grid)
    for channel, mesh, transform in [(0, coarse, coarse_transform), (1, fine, fine_transform)]:
        if mesh:
            target["geometry"][:, channel] = mesh_occupancy(mesh, transform, grid)
            target["geometry_valid"][channel] = 1
    row = {"id": sample_id, "image": f"images/{sample_id}.png", "target": f"targets/{sample_id}.npz",
           "source": source, "grid": grid, "teachers": {"appearance": source,
           "coarse": "InstantMesh" if coarse else None, "fine": "Hi3D" if fine else None},
           "teacher_gaussians": len(splats.xyz), "retained_gaussians": int(target["occ"].sum()),
           "geometry_transforms": {"coarse": json.loads(Path(coarse_transform).read_text()) if coarse else None,
                                   "fine": json.loads(Path(fine_transform).read_text()) if fine else None},
           "metadata": extra or {}}
    manifest = root / "manifest.jsonl"
    rows = load_manifest(manifest) if manifest.exists() else []
    if any(r["id"] == sample_id for r in rows):
        raise ValueError(f"Sample {sample_id} already exists; choose a new ID")
    if rows and any(r["grid"] != grid for r in rows):
        raise ValueError("Do not mix grid resolutions in one dataset")
    image = Image.open(image_path)
    image.load()
    if root.joinpath(row["image"]).exists() or root.joinpath(row["target"]).exists():
        raise ValueError("Sample files already exist; refusing to overwrite")
    root.joinpath("images").mkdir(parents=True, exist_ok=True)
    root.joinpath("targets").mkdir(parents=True, exist_ok=True)
    image.save(root / row["image"])
    row["image_sha256"] = hashlib.sha256((root / row["image"]).read_bytes()).hexdigest()
    np.savez_compressed(root / row["target"], **target)
    row["target_sha256"] = hashlib.sha256((root / row["target"]).read_bytes()).hexdigest()
    # Single-writer dataset. Atomic manifest replacement avoids partial JSON lines.
    tmp = manifest.with_suffix(".tmp")
    tmp.write_text("".join(json.dumps(r) + "\n" for r in [*rows, row]))
    tmp.replace(manifest)
    return row


def load_manifest(path):
    rows = [json.loads(line) for line in Path(path).read_text().splitlines() if line.strip()]
    if not rows or len({r["id"] for r in rows}) != len(rows):
        raise ValueError("Manifest must contain nonempty, unique sample IDs")
    return rows


def image_tensor(path, size=128):
    """Already-framed RGB or RGBA cutout; black background, square letterbox."""
    with Image.open(path) as original:
        rgba = ImageOps.exif_transpose(original).convert("RGBA")
        bg = Image.new("RGBA", rgba.size, (0, 0, 0, 255))
        bg.alpha_composite(rgba)
        image = ImageOps.pad(bg.convert("RGB"), (size, size), method=Image.Resampling.LANCZOS, color="black")
    return torch.from_numpy(np.array(image, dtype=np.float32).transpose(2, 0, 1) / 255)


class CachedDataset(Dataset):
    def __init__(self, manifest, rows=None):
        self.root = Path(manifest).resolve().parent
        self.rows = load_manifest(manifest) if rows is None else rows
        if not self.rows:
            raise ValueError("Empty dataset")
        self.grid = self.rows[0]["grid"]
        for row in self.rows:
            if row["grid"] != self.grid:
                raise ValueError("Mixed grids")
            for key in ("image", "target"):
                p = (self.root / row[key]).resolve()
                if not p.is_relative_to(self.root) or not p.is_file():
                    raise ValueError(f"Invalid dataset {key} path")
                digest = hashlib.sha256(p.read_bytes()).hexdigest()
                if digest != row[f"{key}_sha256"]:
                    raise ValueError(f"Dataset {key} hash mismatch for {row['id']}")

    def __len__(self):
        return len(self.rows)

    def __getitem__(self, i):
        row = self.rows[i]
        with np.load(self.root / row["target"], allow_pickle=False) as cached:
            target = {k: torch.from_numpy(cached[k].copy()).float() for k in cached.files if k != "grid"}
        return image_tensor(self.root / row["image"]), target