from __future__ import annotations

import hashlib
import json
from pathlib import Path
from typing import Any, Mapping

import numpy as np
from PIL import Image


REQUIRED_BANDS = ("green", "red", "red_edge", "nir")


def read_json(path: str | Path) -> dict[str, Any]:
    value = json.loads(Path(path).read_text(encoding="utf-8"))
    if not isinstance(value, dict):
        raise ValueError(f"Expected a JSON object: {path}")
    return value


def write_json(path: str | Path, value: Any) -> Path:
    destination = Path(path)
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_text(
        json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    return destination


def sha256_file(path: str | Path) -> str:
    digest = hashlib.sha256()
    with Path(path).open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def sha256_tree(root: str | Path, pattern: str = "*.py") -> str:
    base = Path(root)
    digest = hashlib.sha256()
    for path in sorted(base.rglob(pattern), key=lambda item: item.relative_to(base).as_posix()):
        relative = path.relative_to(base).as_posix().encode("utf-8")
        content = path.read_bytes()
        digest.update(len(relative).to_bytes(4, "big"))
        digest.update(relative)
        digest.update(len(content).to_bytes(8, "big"))
        digest.update(content)
    return digest.hexdigest()


def save_rgb(path: str | Path, rgb: np.ndarray) -> Path:
    destination = Path(path)
    destination.parent.mkdir(parents=True, exist_ok=True)
    array = np.asarray(rgb, dtype=np.float32)
    if array.ndim != 3 or array.shape[2] != 3:
        raise ValueError("RGB input must have shape (rows, columns, 3)")
    Image.fromarray(np.rint(np.clip(array, 0, 1) * 255).astype(np.uint8), mode="RGB").save(destination)
    return destination


def save_reflectance(path: str | Path, reflectance: np.ndarray, scale_factor: int = 10_000) -> Path:
    destination = Path(path)
    destination.parent.mkdir(parents=True, exist_ok=True)
    array = np.asarray(reflectance, dtype=np.float32)
    if array.ndim != 2:
        raise ValueError("Reflectance input must be a two-dimensional band")
    if scale_factor <= 0:
        raise ValueError("scale_factor must be positive")
    Image.fromarray(np.rint(np.clip(array, 0, 1) * scale_factor).astype(np.uint16)).save(destination)
    return destination


def load_rgb(path: str | Path) -> np.ndarray:
    value = np.asarray(Image.open(path).convert("RGB"), dtype=np.float32) / 255.0
    if value.ndim != 3 or value.shape[2] != 3:
        raise ValueError(f"Expected RGB imagery: {path}")
    return value


def load_reflectance(path: str | Path, scale_factor: float) -> np.ndarray:
    if not np.isfinite(scale_factor) or scale_factor <= 0:
        raise ValueError("scale_factor must be positive")
    value = np.asarray(Image.open(path), dtype=np.float32)
    if value.ndim != 2:
        raise ValueError(f"Expected single-band imagery: {path}")
    return value / float(scale_factor)


def _resolve_manifest_path(manifest_path: Path, value: object) -> Path:
    if not isinstance(value, str) or not value:
        raise ValueError("Capture manifest contains a missing image path")
    candidate = Path(value)
    return candidate if candidate.is_absolute() else manifest_path.parent / candidate


def load_capture(manifest_path: str | Path) -> tuple[np.ndarray, dict[str, np.ndarray], dict[str, Any]]:
    path = Path(manifest_path)
    manifest = read_json(path)
    files = manifest.get("files")
    if not isinstance(files, Mapping):
        raise ValueError("Capture manifest files must be an object")
    resolved = {name: _resolve_manifest_path(path, files.get(name)) for name in ("rgb", *REQUIRED_BANDS)}
    expected_hashes = manifest.get("fileSha256")
    if not isinstance(expected_hashes, Mapping):
        raise ValueError("Capture manifest must provide fileSha256 for every image")
    for name, image_path in resolved.items():
        expected = expected_hashes.get(name)
        if not isinstance(expected, str) or len(expected) != 64 or sha256_file(image_path) != expected:
            raise ValueError(f"Capture image checksum mismatch: {name}")
    rgb = load_rgb(resolved["rgb"])
    scale_factor = float(manifest.get("scaleFactor", 0))
    bands = {
        name: load_reflectance(resolved[name], scale_factor)
        for name in REQUIRED_BANDS
    }
    return rgb, bands, manifest
