"""Memory-mapped checkpoint manifests and atomic publication of masks."""

import fcntl
import json
import os
from pathlib import Path
import shutil
import tempfile

import numpy as np

from .core import _inputs, _options, _run


def _input_path(root, relative, suffix):
    if not isinstance(relative, str) or Path(relative).is_absolute():
        raise ValueError("manifest paths must be relative to the input root")
    path = (root / relative).resolve(strict=True)
    if not path.is_relative_to(root) or not path.is_file() or path.suffix != suffix:
        raise ValueError(f"expected a {suffix} file inside the input root")
    return path


def _open_npy(root, name):
    return np.load(_input_path(root, name, ".npy"), mmap_mode="r", allow_pickle=False)


def _read_manifest(root, manifest, max_weights):
    path = _input_path(root, manifest, ".json")
    with path.open("rb") as file:
        raw = file.read(1_048_577)
    if len(raw) > 1_048_576:
        raise ValueError("manifest exceeds 1 MiB")
    spec = json.loads(raw)
    if not isinstance(spec, dict) or set(spec) != {"keep", "projections"}:
        raise ValueError("manifest must contain keep and projections")
    items = spec["projections"]
    if not isinstance(items, list) or not 1 <= len(items) <= 1024:
        raise ValueError("provide 1 to 1024 projection records")
    if type(spec["keep"]) is not int:
        raise ValueError("manifest keep must be an integer")
    names, weights, moments = [], [], []
    seen = set()
    total = 0
    for item in items:
        if (not isinstance(item, dict) or not {"name", "weights"} <= item.keys()
                or item.keys() - {"name", "weights", "input_second_moment"}):
            raise ValueError("projection records require name and weights")
        name = item["name"]
        if not isinstance(name, str) or not 1 <= len(name) <= 256 or name in seen:
            raise ValueError("projection names must be nonempty and unique")
        seen.add(name)
        w = _open_npy(root, item["weights"])
        total += w.size
        if max_weights is not None and total > max_weights:
            raise ValueError("manifest exceeds the configured weight limit")
        names.append(name)
        weights.append(w)
        moments.append(_open_npy(root, item["input_second_moment"])
                       if "input_second_moment" in item else None)
    return names, weights, moments, spec["keep"]


def prune_manifest(input_root, manifest, output_dir, *, scratch_dir=None,
                   threads=1, seed=0, max_weights=None):
    """Publish disk-backed .npy masks and result.json to a new directory."""
    root = Path(input_root).resolve(strict=True)
    names, weights, moments, keep = _read_manifest(root, manifest, max_weights)
    return _publish_masks(names, weights, moments, keep, output_dir,
                          scratch_dir=scratch_dir, threads=threads, seed=seed)


def _publish_masks(names, weights, moments, keep, output_dir, *,
                   scratch_dir, threads, seed):
    output = Path(output_dir).absolute()
    arrays, moments = _inputs(weights, moments)
    total = sum(w.size for w in arrays)
    keep, threads, seed = _options(total, keep, threads, seed)

    output.parent.mkdir(parents=True, exist_ok=True)
    lock_path = output.parent / f".{output.name}.lock"
    with lock_path.open("a+b") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX)
        if output.exists() or output.is_symlink():
            raise FileExistsError(output)
        stage = Path(tempfile.mkdtemp(prefix=f".{output.name}-", dir=output.parent))
        try:
            paths = [stage / f"mask-{j:04d}.npy" for j in range(len(arrays))]
            masks = [np.lib.format.open_memmap(path, mode="w+", dtype=bool, shape=w.shape)
                     for path, w in zip(paths, arrays)]
            for path in paths:
                with path.open("r+b") as file:
                    os.posix_fallocate(file.fileno(), 0, path.stat().st_size)
            scratch = Path(scratch_dir) if scratch_dir is not None else stage
            _run(arrays, moments, masks, total, keep, threads, seed, scratch)
            entries = []
            for name, path, mask in zip(names, paths, masks):
                mask.flush()
                entries.append(dict(name=name, mask=path.name, shape=list(mask.shape),
                                    kept_slots=int(np.count_nonzero(mask))))
            result = dict(weights=total, kept_slots=keep, projections=entries)
            if sum(item["kept_slots"] for item in entries) != keep:
                raise RuntimeError("native selector returned an incorrect retention budget")
            (stage / "result.json").write_text(json.dumps(result, indent=2) + "\n")
            stage.rename(output)
            return result
        finally:
            if stage.exists():
                shutil.rmtree(stage)