"""I/O utilities: saving depth / edges / segmentation / safetensors.

Directory-based output layout — modality determines the folder, not file suffix:

    output_root/
    ├── depth/{rel_parent}/{stem}.png          # vis
    ├── edge/{rel_parent}/{stem}.png
    ├── segm/{rel_parent}/{stem}.png
    ├── chm/{rel_parent}/{stem}.png
    ├── npy/depth/{rel_parent}/{stem}.npy      # intermediate float16/uint8
    ├── npy/edge/{rel_parent}/{stem}.npy
    ├── npy/segm/{rel_parent}/{stem}.npy
    ├── npy/chm/{rel_parent}/{stem}.npy
    └── safetensors/{rel_parent}/{stem}.safetensors

No global config imports — all parameters passed explicitly.
"""

from __future__ import annotations

import atexit
import logging
import os
import tempfile
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from safetensors.torch import save_file as _st_save_file, load_file as st_load_file

logger = logging.getLogger(__name__)

_palette_cache: dict[int, np.ndarray] = {}

# ---------------------------------------------------------------------------
# Async I/O pool
# ---------------------------------------------------------------------------
_io_pool: ThreadPoolExecutor | None = None
_IO_WORKERS = 4


def get_io_pool() -> ThreadPoolExecutor:
    """Return (lazily created) shared thread pool for async saves."""
    global _io_pool
    if _io_pool is None:
        _io_pool = ThreadPoolExecutor(max_workers=_IO_WORKERS)
        atexit.register(shutdown_io_pool)
    return _io_pool


def shutdown_io_pool() -> None:
    """Wait for all pending writes and shut down the pool."""
    global _io_pool
    if _io_pool is not None:
        _io_pool.shutdown(wait=True)
        _io_pool = None


# ---------------------------------------------------------------------------
# Path helpers
# ---------------------------------------------------------------------------

def vis_path(output_root: Path, modality: str, rel_parent: str, stem: str) -> Path:
    """Build: output_root / modality / rel_parent / stem.png"""
    return output_root / modality / rel_parent / f"{stem}.png"


def npy_path(output_root: Path, modality: str, rel_parent: str, stem: str) -> Path:
    """Build: output_root / npy / modality / rel_parent / stem.npy"""
    return output_root / "npy" / modality / rel_parent / f"{stem}.npy"


def safetensors_path(output_root: Path, rel_parent: str, stem: str) -> Path:
    """Build: output_root / safetensors / rel_parent / stem.safetensors"""
    return output_root / "safetensors" / rel_parent / f"{stem}.safetensors"


# ---------------------------------------------------------------------------
# Palette
# ---------------------------------------------------------------------------

# Intuitive RS segmentation palette: index → RGB.
_FIXED_PALETTE = np.array([
    [0,     0,   0],  # 0:  background   — black
    [220,  40,  40],  # 1:  building     — red
    [160, 160, 160],  # 2:  road         — gray
    [30,  180,  30],  # 3:  vegetation   — green
    [30,  120, 220],  # 4:  water        — blue
    [180, 140,  80],  # 5:  bare ground  — tan
    [120, 100,  80],  # 6:  rock         — brown
    [200, 200,  50],  # 7:  farmland     — yellow
    [100,  60, 120],  # 8:  railway      — purple
    [255, 165,   0],  # 9:  parking lot  — orange
    [200, 200, 200],  # 10: sidewalk     — light gray
    [140, 100,  50],  # 11: bare soil    — dark tan
    [180,  60,  60],  # 12: rooftop      — dark red
    [50,  200, 150],  # 13: sports field — teal
    [80,  100,  70],  # 14: muddy/wetland — olive
    [170, 130, 100],  # 15: embankment  — sandy brown
    [0,   200, 255],  # 16: pool       — cyan
], dtype=np.uint8)


def make_palette(num_classes: int, seed: int = 42) -> np.ndarray:
    """Return color palette for segmentation visualization."""
    if num_classes in _palette_cache:
        return _palette_cache[num_classes]
    if num_classes <= len(_FIXED_PALETTE):
        palette = _FIXED_PALETTE[:num_classes].copy()
    else:
        rng = np.random.RandomState(seed)
        palette = rng.randint(0, 255, (num_classes, 3), dtype=np.uint8)
        palette[:len(_FIXED_PALETTE)] = _FIXED_PALETTE
    _palette_cache[num_classes] = palette
    return palette


# ---------------------------------------------------------------------------
# Low-level atomic save
# ---------------------------------------------------------------------------

def _atomic_save_npy(arr: np.ndarray, path: Path) -> None:
    """Write .npy atomically via temp file + rename."""
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp = tempfile.mkstemp(suffix=".npy", dir=path.parent)
    os.close(fd)
    try:
        np.save(tmp, arr)
        os.replace(tmp, path)
    except BaseException:
        if os.path.exists(tmp):
            os.remove(tmp)
        raise


_COLORMAP_CACHE: dict[str, np.ndarray] = {}


def _apply_colormap(gray: np.ndarray, cmap_name: str = "turbo") -> np.ndarray:
    """Apply matplotlib colormap to [H, W] float32 [0, 1] → [H, W, 3] uint8."""
    if cmap_name not in _COLORMAP_CACHE:
        import matplotlib.cm as cm
        cmap = cm.get_cmap(cmap_name)
        lut = (cmap(np.linspace(0, 1, 256))[:, :3] * 255).astype(np.uint8)
        _COLORMAP_CACHE[cmap_name] = lut
    lut = _COLORMAP_CACHE[cmap_name]
    idx = (gray.clip(0, 1) * 255).astype(np.uint8)
    return lut[idx]


# ---------------------------------------------------------------------------
# Save float16 maps (depth, edge, chm)
# ---------------------------------------------------------------------------

def _save_float16_map(
    data: torch.Tensor,
    output_root: Path,
    rel_parent: str,
    stem: str,
    modality: str,
    save_npy: bool = True,
    save_vis: bool = True,
    colormap: str | None = None,
) -> None:
    """Save a [1, H, W] float tensor as .npy (float16) + optional vis .png."""
    arr = data.half().numpy()
    if save_npy:
        p = npy_path(output_root, modality, rel_parent, stem)
        _atomic_save_npy(arr, p)
    if save_vis:
        gray = arr.squeeze(0).astype(np.float32)
        if colormap:
            vis = _apply_colormap(gray, colormap)
        else:
            vis = (gray * 255).clip(0, 255).astype(np.uint8)
        p = vis_path(output_root, modality, rel_parent, stem)
        p.parent.mkdir(parents=True, exist_ok=True)
        Image.fromarray(vis).save(p)


def save_depth(depth: torch.Tensor, output_root: Path, rel_parent: str,
               stem: str, save_npy: bool = True, save_vis: bool = True) -> None:
    _save_float16_map(depth, output_root, rel_parent, stem, "depth", save_npy, save_vis)


def save_depth_async(depth: torch.Tensor, output_root: Path, rel_parent: str,
                     stem: str, save_npy: bool = True, save_vis: bool = True) -> None:
    get_io_pool().submit(save_depth, depth.clone().cpu(), output_root, rel_parent,
                         stem, save_npy, save_vis)


def save_chmv2(depth: torch.Tensor, output_root: Path, rel_parent: str,
               stem: str, save_npy: bool = True, save_vis: bool = True) -> None:
    _save_float16_map(depth, output_root, rel_parent, stem, "chm", save_npy, save_vis)


def save_chmv2_async(depth: torch.Tensor, output_root: Path, rel_parent: str,
                     stem: str, save_npy: bool = True, save_vis: bool = True) -> None:
    get_io_pool().submit(save_chmv2, depth.clone().cpu(), output_root, rel_parent,
                         stem, save_npy, save_vis)


def save_edges(edges: torch.Tensor, output_root: Path, rel_parent: str,
               stem: str, save_npy: bool = True, save_vis: bool = True) -> None:
    _save_float16_map(edges, output_root, rel_parent, stem, "edge", save_npy, save_vis)


def save_edges_async(edges: torch.Tensor, output_root: Path, rel_parent: str,
                     stem: str, save_npy: bool = True, save_vis: bool = True) -> None:
    get_io_pool().submit(save_edges, edges.clone().cpu(), output_root, rel_parent,
                         stem, save_npy, save_vis)


# ---------------------------------------------------------------------------
# Save segmentation
# ---------------------------------------------------------------------------

def save_segmentation(
    seg_ids: torch.Tensor,
    output_root: Path,
    rel_parent: str,
    stem: str,
    save_npy: bool = True,
    save_vis: bool = True,
    num_classes: int = 150,
) -> None:
    """Save segmentation map [1, H, W] uint8."""
    arr = seg_ids.byte().numpy()
    if save_npy:
        _atomic_save_npy(arr, npy_path(output_root, "segm", rel_parent, stem))
    if save_vis:
        palette = make_palette(num_classes)
        seg_np = arr.squeeze(0).astype(np.uint8)
        seg_clamped = np.clip(seg_np, 0, num_classes - 1).astype(np.uint8)
        img = Image.fromarray(seg_clamped).convert("P")
        flat_pal = np.zeros(768, dtype=np.uint8)
        flat_pal[: num_classes * 3] = palette.flatten()
        img.putpalette(flat_pal.tolist())
        p = vis_path(output_root, "segm", rel_parent, stem)
        p.parent.mkdir(parents=True, exist_ok=True)
        img.save(p)


def save_segmentation_async(
    seg_ids: torch.Tensor,
    output_root: Path,
    rel_parent: str,
    stem: str,
    save_npy: bool = True,
    save_vis: bool = True,
    num_classes: int = 150,
) -> None:
    get_io_pool().submit(
        save_segmentation, seg_ids.clone().cpu(), output_root, rel_parent,
        stem, save_npy, save_vis, num_classes,
    )


# ---------------------------------------------------------------------------
# SafeTensors: consolidate all modalities into one file per image
# ---------------------------------------------------------------------------

_MODALITY_SPEC: dict[str, tuple[torch.dtype, str]] = {
    "depth": (torch.float16, "depth"),
    "edge":  (torch.float16, "edge"),
    "chm":   (torch.float16, "chm"),
    "segm":  (torch.uint8,   "segm"),
}


def _load_modality_tensor(
    output_root: Path, rel_parent: str, stem: str,
    modality: str, dtype: torch.dtype,
) -> torch.Tensor | None:
    """Load a single modality from .npy or .png, return [1, H, W] tensor or None."""
    np_p = npy_path(output_root, modality, rel_parent, stem)
    vis_p = vis_path(output_root, modality, rel_parent, stem)

    if np_p.exists():
        arr = np.load(np_p)
        t = torch.from_numpy(arr.astype(np.float32 if dtype != torch.uint8 else np.uint8))
        if t.ndim == 2:
            t = t.unsqueeze(0)
        return t.to(dtype)

    if vis_p.exists():
        if modality == "segm":
            pil = Image.open(vis_p)
            if pil.mode == "P":
                img = np.array(pil)
            else:
                logger.debug("Skipping %s segm.png (RGB, no class IDs).", stem)
                return None
            t = torch.from_numpy(img.astype(np.uint8))
            if t.ndim == 2:
                t = t.unsqueeze(0)
            return t
        else:
            img = np.array(Image.open(vis_p))
            arr = img.astype(np.float32) / 255.0
            if arr.ndim == 2:
                arr = arr[np.newaxis]
            elif arr.ndim == 3:
                arr = arr[:, :, 0:1].transpose(2, 0, 1)
            return torch.from_numpy(arr).to(dtype)

    return None


def consolidate_safetensors(
    output_root: Path,
    rel_parent: str,
    stem: str,
    cleanup_npy: bool = False,
) -> bool:
    """Bundle available modalities into one .safetensors file.

    Returns True if the file was written, False if no modalities found.
    """
    tensors: dict[str, torch.Tensor] = {}
    npy_paths_to_clean: list[Path] = []

    for modality, (dtype, _) in _MODALITY_SPEC.items():
        t = _load_modality_tensor(output_root, rel_parent, stem, modality, dtype)
        if t is not None:
            tensors[modality] = t
            np_p = npy_path(output_root, modality, rel_parent, stem)
            if np_p.exists():
                npy_paths_to_clean.append(np_p)

    if not tensors:
        return False

    st_p = safetensors_path(output_root, rel_parent, stem)
    st_p.parent.mkdir(parents=True, exist_ok=True)

    fd, tmp = tempfile.mkstemp(suffix=".safetensors", dir=st_p.parent)
    os.close(fd)
    try:
        _st_save_file(tensors, tmp)
        os.replace(tmp, st_p)
    except BaseException:
        if os.path.exists(tmp):
            os.remove(tmp)
        raise

    if cleanup_npy:
        for p in npy_paths_to_clean:
            p.unlink(missing_ok=True)

    return True


def consolidate_safetensors_async(
    output_root: Path,
    rel_parent: str,
    stem: str,
    cleanup_npy: bool = False,
) -> None:
    get_io_pool().submit(consolidate_safetensors, output_root, rel_parent,
                         stem, cleanup_npy)


# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------

def setup_logging(log_level: str = "INFO", log_file: Path | None = None) -> None:
    """Configure root logger with coloredlogs for console + optional file handler."""
    import coloredlogs

    fmt = "%(asctime)s | %(levelname)-7s | %(name)s | %(message)s"
    datefmt = "%H:%M:%S"
    level = getattr(logging, log_level)

    coloredlogs.install(
        level=level,
        fmt=fmt,
        datefmt=datefmt,
        level_styles={
            "debug": {"color": "cyan"},
            "info": {"color": "green"},
            "warning": {"color": "yellow", "bold": True},
            "error": {"color": "red", "bold": True},
            "critical": {"color": "red", "bold": True, "background": "white"},
        },
        field_styles={
            "asctime": {"color": "white", "faint": True},
            "levelname": {"color": "magenta", "bold": True},
            "name": {"color": "blue"},
        },
    )

    if log_file is not None:
        log_file.parent.mkdir(parents=True, exist_ok=True)
        file_handler = logging.FileHandler(log_file, encoding="utf-8")
        file_handler.setFormatter(logging.Formatter(fmt, datefmt=datefmt))
        logging.root.addHandler(file_handler)
