"""Utilities for loading and iterating the S23DR 2026 dataset.

IMPORTANT: The dataset contains at least one corrupt image (around row ~13077
in the training split). All iteration helpers in this module wrap entries in
try/except so a single bad scene never kills the entire run.
"""

import io
import os
import tempfile
import zipfile
import warnings
from typing import Dict, List, Optional, Tuple

import numpy as np
from PIL import Image, UnidentifiedImageError


# ---------------------------------------------------------------------------
# Safe entry validation
# ---------------------------------------------------------------------------

def validate_entry(entry: dict) -> bool:
    """Check whether a dataset entry has valid (non-corrupt) images.

    Returns False if any image column contains an unreadable PIL image.
    This catches the PIL.UnidentifiedImageError that occurs around row ~13077.
    """
    try:
        for col in ("gestalt", "ade", "depth"):
            imgs = entry.get(col, [])
            if imgs is None:
                return False
            for img in imgs:
                if img is None:
                    return False
                # Force PIL to actually decode the pixels
                np.array(img)
        return True
    except (UnidentifiedImageError, OSError, ValueError, Exception):
        return False


# ---------------------------------------------------------------------------
# Dataset loading
# ---------------------------------------------------------------------------

def load_dataset_streaming(split: str = "train", token: Optional[str] = None):
    """Load the dataset in streaming mode (low memory, no download)."""
    from datasets import load_dataset
    if token is None:
        token = os.environ.get("HF_TOKEN")
    ds = load_dataset(
        "usm3d/hoho22k_2026_trainval",
        streaming=True,
        trust_remote_code=True,
        token=token,
    )
    return ds[split]


def load_dataset_cached(split: str = "train", cache_dir: str = "/data/s23dr",
                        token: Optional[str] = None):
    """Load dataset with caching to disk for faster repeat access."""
    from datasets import load_dataset
    if token is None:
        token = os.environ.get("HF_TOKEN")
    ds = load_dataset(
        "usm3d/hoho22k_2026_trainval",
        split=split,
        cache_dir=cache_dir,
        trust_remote_code=True,
        token=token,
    )
    return ds


def download_dataset(cache_dir: str = "/data/s23dr", token: Optional[str] = None):
    """Download the full dataset to disk (both train + validation).

    Usage on Modal::

        modal run modal_download_data.py

    Or locally::

        python -c "from src.data_loader import download_dataset; download_dataset('./data')"
    """
    from datasets import load_dataset
    if token is None:
        token = os.environ.get("HF_TOKEN")
    print(f"Downloading train split to {cache_dir} ...")
    load_dataset(
        "usm3d/hoho22k_2026_trainval",
        split="train",
        cache_dir=cache_dir,
        trust_remote_code=True,
        token=token,
    )
    print(f"Downloading validation split to {cache_dir} ...")
    load_dataset(
        "usm3d/hoho22k_2026_trainval",
        split="validation",
        cache_dir=cache_dir,
        trust_remote_code=True,
        token=token,
    )
    print("✅ Dataset downloaded.")


# ---------------------------------------------------------------------------
# Safe iterators  (skip corrupt entries automatically)
# ---------------------------------------------------------------------------

def iter_entries_safe(dataset, max_entries: Optional[int] = None,
                      verbose: bool = True):
    """Iterate over dataset entries, **skipping corrupt rows**.

    Every entry is wrapped in try/except so a single broken image (e.g. the
    known PIL.UnidentifiedImageError near row ~13077) does not crash the run.

    Yields
    ------
    idx : int
        The original row number (before skipping).
    entry : dict
        A validated dataset row.
    """
    from tqdm import tqdm

    yielded = 0
    skipped = 0
    for idx, entry in enumerate(tqdm(dataset, total=max_entries,
                                     desc="Processing entries")):
        if max_entries and yielded >= max_entries:
            break
        try:
            if not validate_entry(entry):
                skipped += 1
                if verbose:
                    oid = entry.get("order_id", f"row_{idx}")
                    print(f"  ⚠ Skipping corrupt entry {oid} (row {idx})")
                continue
            yield idx, entry
            yielded += 1
        except (UnidentifiedImageError, OSError, Exception) as exc:
            skipped += 1
            if verbose:
                print(f"  ⚠ Skipping row {idx}: {type(exc).__name__}: {exc}")

    if verbose and skipped:
        print(f"  ℹ Skipped {skipped} corrupt entries total")


def iter_entries(dataset, max_entries: Optional[int] = None):
    """Legacy helper — yields entries only (no index)."""
    for _, entry in iter_entries_safe(dataset, max_entries=max_entries,
                                      verbose=False):
        yield entry


# ---------------------------------------------------------------------------
# COLMAP helpers
# ---------------------------------------------------------------------------

def read_colmap_reconstruction(colmap_bytes: bytes):
    import pycolmap
    with tempfile.TemporaryDirectory() as tmpdir:
        with zipfile.ZipFile(io.BytesIO(colmap_bytes), "r") as zf:
            zf.extractall(tmpdir)
        rec = pycolmap.Reconstruction(tmpdir)
    return rec


def get_colmap_points(rec) -> np.ndarray:
    pts = [p3d.xyz for p3d in rec.points3D.values()]
    return np.array(pts) if pts else np.zeros((0, 3))


def get_colmap_points_with_colors(rec) -> Tuple[np.ndarray, np.ndarray]:
    pts, cols = [], []
    for p3d in rec.points3D.values():
        pts.append(p3d.xyz)
        cols.append(p3d.color)
    if not pts:
        return np.zeros((0, 3)), np.zeros((0, 3), dtype=np.uint8)
    return np.array(pts), np.array(cols, dtype=np.uint8)


# ---------------------------------------------------------------------------
# Camera / projection helpers
# ---------------------------------------------------------------------------

def depth_to_meters(depth_image) -> np.ndarray:
    return np.array(depth_image).astype(np.float32) / 1000.0


def get_valid_camera_indices(entry: dict) -> List[int]:
    flags = entry.get("pose_only_in_colmap", [])
    return [i for i, f in enumerate(flags) if not f]


def get_camera_params(entry: dict, cam_idx: int):
    K = np.array(entry["K"][cam_idx])
    R = np.array(entry["R"][cam_idx])
    t = np.array(entry["t"][cam_idx])
    return K, R, t


def project_3d_to_2d(points_3d, K, R, t):
    t = t.reshape(3, 1) if t.ndim == 1 else t
    pts_cam = R @ points_3d.T + t
    depth = pts_cam[2, :]
    pts_2d = K @ pts_cam
    uv = pts_2d[:2, :] / pts_2d[2:3, :]
    return uv.T, depth


def backproject_2d_to_3d(uv, depth, K, R, t):
    t = t.reshape(3, 1) if t.ndim == 1 else t
    fx, fy = K[0, 0], K[1, 1]
    cx, cy = K[0, 2], K[1, 2]
    x_n = (uv[:, 0] - cx) / fx
    y_n = (uv[:, 1] - cy) / fy
    pts_cam = np.stack([x_n * depth, y_n * depth, depth], axis=1)
    R_inv = R.T
    t_inv = -R.T @ t
    return (R_inv @ pts_cam.T + t_inv).T