"""
Shared driver for the upstream NV-Generate-CTMR scripts.

Strategy: the upstream code is structured around argparse + global filesystem layout
(reads relative-path configs, writes to a relative output_dir). Rather than refactor
its internals, we treat it as an external tool: chdir into the upstream root, write
modified config copies that override the user-controlled fields, ensure weights are
downloaded, then call the upstream entry-point function. We then return the path of
the most recently produced NIfTI from the configured output dir.
"""
from __future__ import annotations

import contextlib
import importlib
from collections import OrderedDict
import json
import os
import sys
import time
import uuid
from pathlib import Path
from typing import Iterable, Optional

from utils.timing import Stopwatch, log, phase

ROOT = Path(__file__).resolve().parent.parent
UPSTREAM = ROOT / "repos" / "NV-Generate-CTMR"
GENERATED_OUTPUT = UPSTREAM / "output"

# Upstream `scripts.inference` reads MONAI_DATA_DIRECTORY and passes it to
# download_model_data() as root_dir, which is where it expects the `datasets/`
# payload to live. We prefetch to the same place so it finds the files already
# there instead of downloading a second copy.
DATA_DIR = UPSTREAM / "temp_work_dir"

MASK_ARCHIVE = DATA_DIR / "datasets" / "all_masks_flexible_size_and_spacing_4000.zip"

ALL_VERSIONS = ("rflow-ct", "rflow-mr", "rflow-mr-brain")

# Once the mask archive is extracted, nothing reads the .zip again -- but
# upstream's download_model_data() re-fetches any file whose destination is
# missing, and that is an 11.6 GB round trip. Truncating it to a placeholder
# keeps the path present for that existence check while giving back the space.
# prefetch_weights() repairs the state if the extracted directory ever goes
# away. Set NV_GENERATE_PRUNE_ARCHIVE=0 to keep the full archive on disk.
PRUNE_ARCHIVE = os.environ.get("NV_GENERATE_PRUNE_ARCHIVE", "1") == "1"

# Cap for the checkpoint cache below, in GiB of host RAM.
CHECKPOINT_CACHE_GIB = float(os.environ.get("NV_GENERATE_CKPT_CACHE_GIB", "8"))


class _LinkingShutil:
    """Drop-in for the `shutil` module that hard-links instead of copying.

    Upstream's fetch_to_hf_path_cmd() calls hf_hub_download() (which stores the
    blob in the HF cache) and then shutil.copy2(cached, dst) — so every file is
    written to disk twice. For CT that is 11.6 GB of anatomy masks duplicated.
    A hard link produces the same destination path for zero extra bytes.

    Moving the blob out of the cache instead is not safe: hf_hub_download keeps
    a blobs/ + snapshots/ layout and removing the blob both corrupts the cache
    and forces a full re-download on the next call. A link keeps one inode with
    two names, so either path can be removed independently.

    Falls back to a symlink, then to a real copy, if the filesystem refuses
    (e.g. the cache and the app are on different mounts).
    """

    def __init__(self, real):
        self._real = real

    def __getattr__(self, name):
        return getattr(self._real, name)

    def copy2(self, src, dst, *args, **kwargs):
        # Resolve first. hf_hub_download returns the snapshot entry, which is a
        # *relative* symlink into blobs/ (../../blobs/<sha>). os.link() on a
        # symlink source copies the link itself, target string and all, so
        # linking it into <upstream>/models/ would produce a symlink resolving
        # to <upstream>/../blobs/<sha> — which does not exist. Link the blob.
        src_real = os.path.realpath(os.fspath(src))
        dst_s = os.fspath(dst)

        # Clear any stale entry. A dangling symlink is invisible to exists() but
        # makes link()/symlink() fail EEXIST and makes copy2 write *through* it.
        if os.path.lexists(dst_s):
            try:
                os.unlink(dst_s)
            except OSError:
                pass

        for attempt in (
            lambda: os.link(src_real, dst_s),
            lambda: os.symlink(src_real, dst_s),
        ):
            try:
                attempt()
            except OSError:
                continue
            if os.path.exists(dst_s):  # must actually resolve to real bytes
                return dst_s
            try:
                os.unlink(dst_s)
            except OSError:
                pass
        return self._real.copy2(src_real, dst_s, *args, **kwargs)


@contextlib.contextmanager
def _linking_downloads(download_mod):
    """Make upstream downloads link rather than copy, for the duration of the block.

    Swaps the module's own `shutil` binding rather than patching shutil.copy2
    globally, so nothing outside the upstream download path is affected.
    """
    real = download_mod.shutil
    download_mod.shutil = _LinkingShutil(real)
    try:
        yield
    finally:
        download_mod.shutil = real


def _download_module():
    return importlib.import_module("scripts.download_model_data")


# key -> (payload, nbytes). Sizes are tracked here rather than as an attribute
# on the payload, because state dicts are plain dicts and reject attributes.
_ckpt_cache: "OrderedDict[tuple[str, int, int], tuple[object, int]]" = OrderedDict()
_ckpt_bytes = 0


@contextlib.contextmanager
def _cached_torch_load():
    """Serve repeat torch.load() calls for the same checkpoint from RAM.

    scripts/inference.py performs five torch.load() calls per CT run -- the
    autoencoder, diffusion UNet, controlnet, and both mask-generation nets,
    ~3.35 GB in total -- and re-reads every one of them from disk on each
    generation. The deserialised objects are only read from (load_state_dict
    copies out of them), so handing back the same object is safe.

    This caches host-side state dicts, not GPU memory: weights are ~3.35 GB for
    CT against a 96 GB xlarge slice, so VRAM was never the limit here. RAM is,
    hence the byte budget and LRU eviction.
    """
    import torch

    original = torch.load

    def _load(f, *args, **kwargs):
        global _ckpt_bytes
        try:
            key_path = os.path.realpath(os.fspath(f))
            st = os.stat(key_path)
            key = (key_path, st.st_size, int(st.st_mtime))
        except (TypeError, OSError):
            return original(f, *args, **kwargs)  # file-like object, not a path

        if key in _ckpt_cache:
            _ckpt_cache.move_to_end(key)
            log(f"   ckpt cache HIT  {os.path.basename(key_path)}")
            return _ckpt_cache[key][0]

        obj = original(f, *args, **kwargs)
        budget = int(CHECKPOINT_CACHE_GIB * 1024 ** 3)
        if st.st_size <= budget:
            while _ckpt_bytes + st.st_size > budget and _ckpt_cache:
                _, (_evicted, nbytes) = _ckpt_cache.popitem(last=False)
                _ckpt_bytes -= nbytes
                del _evicted
            _ckpt_cache[key] = (obj, st.st_size)
            _ckpt_bytes += st.st_size
            log(f"   ckpt cache MISS {os.path.basename(key_path)}"
                f" (cached {_ckpt_bytes / 1e9:.2f} GB)")
        return obj

    torch.load = _load
    try:
        yield
    finally:
        torch.load = original


@contextlib.contextmanager
def upstream_context():
    """Temporarily add the upstream repo to sys.path and switch CWD to it."""
    if not UPSTREAM.exists():
        raise RuntimeError(
            f"Upstream repo not found at {UPSTREAM}. Run `bash pre-build.sh` first."
        )
    prev_cwd = os.getcwd()
    added = False
    try:
        upstream_str = str(UPSTREAM)
        if upstream_str not in sys.path:
            sys.path.insert(0, upstream_str)
            added = True
        os.chdir(upstream_str)
        yield UPSTREAM
    finally:
        os.chdir(prev_cwd)
        if added and upstream_str in sys.path:
            sys.path.remove(upstream_str)


def ensure_weights(version: str, root_dir: str | None = None) -> None:
    """Download model weights from HF Hub if not already on disk.

    `root_dir` only affects the `datasets/` payload; model checkpoints always
    land in <upstream>/models. Defaults to DATA_DIR so the CT datasets end up
    where upstream `scripts.inference` looks for them.
    """
    if root_dir is None:
        root_dir = str(DATA_DIR)
    with upstream_context():
        download_mod = _download_module()
        # download_model_data is idempotent — it skips files that already exist.
        model_only = version not in ("rflow-ct", "ddpm-ct")
        with _linking_downloads(download_mod):
            download_mod.download_model_data(version, root_dir, model_only=model_only)


def extract_mask_archive() -> bool:
    """Unpack the CT anatomy-mask archive if it is not already unpacked.

    Upstream does this lazily in scripts/find_masks.py on the first CT run,
    which cost ~100s inside the GPU window. Doing it here means find_masks
    finds the directory present and skips straight past. Returns True if this
    call did the extraction.
    """
    mask_dir = MASK_ARCHIVE.parent / MASK_ARCHIVE.stem
    if mask_dir.exists() or not MASK_ARCHIVE.is_file():
        return False
    if MASK_ARCHIVE.stat().st_size == 0:
        return False  # pruned placeholder; repair_mask_archive_state handles it
    from monai.apps.utils import extractall

    extractall(filepath=str(MASK_ARCHIVE), output_dir=str(MASK_ARCHIVE.parent), file_type="zip")
    return True


def repair_mask_archive_state() -> None:
    """Undo a prune whose extracted directory has since disappeared.

    Removing the placeholder makes the destination look absent again, so the
    next download_model_data() refetches the real archive rather than trying to
    unzip an empty file.
    """
    mask_dir = MASK_ARCHIVE.parent / MASK_ARCHIVE.stem
    if MASK_ARCHIVE.is_file() and MASK_ARCHIVE.stat().st_size == 0 and not mask_dir.exists():
        log("mask archive placeholder without extracted data; forcing refetch")
        MASK_ARCHIVE.unlink()


def prune_mask_archive() -> int:
    """Truncate the extracted archive to a placeholder. Returns bytes freed."""
    mask_dir = MASK_ARCHIVE.parent / MASK_ARCHIVE.stem
    if not (PRUNE_ARCHIVE and mask_dir.is_dir() and MASK_ARCHIVE.is_file()):
        return 0
    if not any(mask_dir.iterdir()):
        return 0  # extraction produced nothing; keep the archive
    freed = MASK_ARCHIVE.stat().st_size
    if freed == 0:
        return 0
    MASK_ARCHIVE.unlink()           # drop our link; the cache blob goes too if last
    MASK_ARCHIVE.touch()            # placeholder keeps the existence check happy
    return freed


def prefetch_weights(versions: Iterable[str] = ALL_VERSIONS) -> tuple[Stopwatch, list[str]]:
    """Warm the weight cache at startup, outside any @spaces.GPU function.

    ~15 GB for CT (11.6 GB of that the anatomy-mask archive) plus ~2.3 GB each
    for MR and MR-Brain. Fetching this lazily inside the decorated function
    burned the whole ZeroGPU allocation on I/O before inference began. Failures
    are logged and swallowed: the lazy path still works, it is just slow.
    """
    os.environ.setdefault("MONAI_DATA_DIRECTORY", str(DATA_DIR))
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    watch = Stopwatch("prefetch")
    failures: list[str] = []
    repair_mask_archive_state()
    for version in versions:
        with watch.split(f"fetch {version}"):
            try:
                log(f"-> fetching weights: {version}")
                ensure_weights(version)
            except Exception as e:  # noqa: BLE001 - startup must not be fatal
                failures.append(version)
                log(f"xx fetch FAILED for {version}: {type(e).__name__}: {e}")
    with watch.split("extract mask archive"):
        try:
            if extract_mask_archive():
                log("<- extracted CT anatomy-mask archive")
            else:
                log("<- mask archive already extracted, skipped")
        except Exception as e:  # noqa: BLE001
            failures.append("mask-archive")
            log(f"xx mask archive extraction FAILED: {type(e).__name__}: {e}")
    with watch.split("prune mask archive"):
        try:
            freed = prune_mask_archive()
            if freed:
                log(f"<- reclaimed {freed / 1e9:.1f} GB from the extracted archive")
        except Exception as e:  # noqa: BLE001
            log(f"xx archive prune skipped: {type(e).__name__}: {e}")
    return watch, failures


def _list_outputs_before(output_dir: Path) -> set[str]:
    if not output_dir.exists():
        return set()
    return {p.name for p in output_dir.glob("*.nii.gz")}


def _newest_outputs(output_dir: Path, before: set[str]) -> list[Path]:
    if not output_dir.exists():
        return []
    new = [p for p in output_dir.glob("*.nii.gz") if p.name not in before]
    new.sort(key=lambda p: p.stat().st_mtime)
    return new


def _write_temp_configs(
    base_env_config: Path,
    base_model_config: Path,
    overrides: dict,
    tag: str,
) -> tuple[Path, Path]:
    """
    Write modified copies of the env + model configs into a per-call temp dir under
    UPSTREAM / configs / _temp /.  Returns (env_path, model_path).
    """
    temp_dir = UPSTREAM / "configs" / "_temp"
    temp_dir.mkdir(parents=True, exist_ok=True)

    env = json.loads(base_env_config.read_text())
    model = json.loads(base_model_config.read_text())
    if "env" in overrides:
        env.update(overrides["env"])
    if "diffusion_unet_inference" in overrides:
        model.setdefault("diffusion_unet_inference", {}).update(overrides["diffusion_unet_inference"])

    suffix = f"{tag}_{uuid.uuid4().hex[:8]}"
    env_path = temp_dir / f"env_{suffix}.json"
    model_path = temp_dir / f"model_{suffix}.json"
    env_path.write_text(json.dumps(env, indent=2))
    model_path.write_text(json.dumps(model, indent=2))
    return env_path, model_path


def run_image_only(
    *,
    version: str,
    output_size: tuple[int, int, int],
    spacing: tuple[float, float, float],
    modality: int,
    seed: int,
    num_inference_steps: int = 30,
    cfg_guidance_scale: Optional[float] = None,
) -> Path:
    """
    Run the image-only diffusion pipeline (`scripts.diff_model_infer`) for the given
    version (rflow-ct / rflow-mr / rflow-mr-brain). Returns path to generated NIfTI.
    """
    ensure_weights(version)

    base_env = UPSTREAM / "configs" / f"environment_maisi_diff_model_{version}.json"
    base_model = UPSTREAM / "configs" / f"config_maisi_diff_model_{version}.json"
    network_def = UPSTREAM / "configs" / "config_network_rflow.json"

    inference_overrides = {
        "dim": list(output_size),
        "spacing": list(spacing),
        "modality": modality,
        "random_seed": seed,
        "num_inference_steps": num_inference_steps,
    }
    if cfg_guidance_scale is not None:
        inference_overrides["cfg_guidance_scale"] = cfg_guidance_scale

    with upstream_context():
        env_path, model_path = _write_temp_configs(
            base_env_config=base_env,
            base_model_config=base_model,
            overrides={"diffusion_unet_inference": inference_overrides},
            tag=version,
        )

        # Read env to determine output_dir (relative to upstream root)
        env_data = json.loads(env_path.read_text())
        output_dir = (UPSTREAM / env_data["output_dir"]).resolve()
        existing = _list_outputs_before(output_dir)

        diff_mod = importlib.import_module("scripts.diff_model_infer")

        t0 = time.time()
        with phase(f"{version}.diffusion"), _cached_torch_load():
            diff_mod.diff_model_infer(
                env_config_path=str(env_path.relative_to(UPSTREAM)),
                model_config_path=str(model_path.relative_to(UPSTREAM)),
                model_def_path=str(network_def.relative_to(UPSTREAM)),
                num_gpus=1,
            )
        runtime = time.time() - t0

        new_files = _newest_outputs(output_dir, existing)
        if not new_files:
            raise RuntimeError(f"No new NIfTI produced in {output_dir}")
        latest = new_files[-1]

    # Cleanup temp configs (don't fail if cleanup errors)
    for p in (env_path, model_path):
        try:
            p.unlink()
        except OSError:
            pass

    return latest


def run_paired_ct(
    *,
    output_size: tuple[int, int, int],
    spacing: tuple[float, float, float],
    body_region: list[str],
    anatomy_list: list[str],
    seed: int,
    num_inference_steps: int = 30,
    num_output_samples: int = 1,
) -> tuple[Path, Optional[Path]]:
    """
    Run the paired CT image+mask pipeline (`scripts.inference`). Returns
    (image_path, mask_path). Mask is the corresponding label volume.
    """
    version = "rflow-ct"
    # No ensure_weights() here: upstream scripts/inference.py calls
    # download_model_data(version, MONAI_DATA_DIRECTORY) itself. Calling it here
    # too — with a different root_dir — downloaded the 11.6 GB anatomy-mask
    # archive to a second location. prefetch_weights() warms the same paths at
    # startup, so this run finds everything already present.

    base_env = UPSTREAM / "configs" / f"environment_{version}.json"
    base_infer = UPSTREAM / "configs" / "config_infer.json"

    # Build a custom config_infer with overrides
    infer_data = json.loads(base_infer.read_text())
    infer_data["output_size"] = list(output_size)
    infer_data["spacing"] = list(spacing)
    infer_data["body_region"] = list(body_region)
    infer_data["anatomy_list"] = list(anatomy_list)
    infer_data["num_inference_steps"] = num_inference_steps
    infer_data["num_output_samples"] = num_output_samples

    temp_dir = UPSTREAM / "configs" / "_temp"
    temp_dir.mkdir(parents=True, exist_ok=True)
    suffix = uuid.uuid4().hex[:8]
    infer_path = temp_dir / f"config_infer_{version}_{suffix}.json"
    infer_path.write_text(json.dumps(infer_data, indent=2))

    env_data = json.loads(base_env.read_text())
    output_dir = (UPSTREAM / env_data["output_dir"]).resolve()

    with upstream_context():
        existing = _list_outputs_before(output_dir)

        inference_mod = importlib.import_module("scripts.inference")
        # The upstream `main()` parses argv directly. Patch sys.argv around the call.
        old_argv = sys.argv
        sys.argv = [
            "scripts.inference",
            "-t", "./configs/config_network_rflow.json",
            "-i", str(infer_path.relative_to(UPSTREAM)),
            "-e", str(base_env.relative_to(UPSTREAM)),
            "--random-seed", str(seed),
            "--version", version,
        ]
        os.environ.setdefault("MONAI_DATA_DIRECTORY", str(DATA_DIR))

        try:
            # Link-instead-of-copy also covers the download inference.py runs
            # internally, for the case where prefetch was skipped or failed.
            with phase("ct.paired-inference"), _cached_torch_load(), _linking_downloads(_download_module()):
                inference_mod.main()
        finally:
            sys.argv = old_argv

        new_files = _newest_outputs(output_dir, existing)

    try:
        infer_path.unlink()
    except OSError:
        pass

    # Paired pipeline writes both image and label NIfTIs. Convention: filenames
    # contain "image" / "label" or are emitted as adjacent files.
    image_path: Optional[Path] = None
    mask_path: Optional[Path] = None
    for p in new_files:
        name = p.name.lower()
        if "label" in name or "_mask" in name or "seg" in name:
            mask_path = p
        elif "image" in name or "img" in name:
            image_path = p

    # Fallback: if naming is ambiguous, treat the smaller-modality-time file as image
    if image_path is None and new_files:
        image_path = new_files[0]
    if mask_path is None and len(new_files) > 1:
        mask_path = new_files[-1]

    if image_path is None:
        raise RuntimeError(f"No NIfTI produced in {output_dir}")
    return image_path, mask_path


def labels_present(mask_path: Path) -> set[int]:
    """Return the set of unique non-zero label IDs present in the mask volume."""
    import nibabel as nib
    import numpy as np

    with phase("ct.label-scan"):
        img = nib.load(str(mask_path))
        data = np.asarray(img.dataobj)
        uniq = np.unique(data).astype(int).tolist()
        return {int(u) for u in uniq if u != 0}