#!/usr/bin/env python3
"""
setup.py - One-command setup for the CD-Models benchmark suite.

Clones all model repositories, downloads pretrained backbone weights,
verifies the Python environment, and optionally prepares datasets.

Usage:
    python setup.py                     # full setup (clone + weights + env check)
    python setup.py --skip-weights      # clone only, skip weight downloads
    python setup.py --env-check-only    # only check Python packages
    python setup.py --status            # print current setup status without changes
    python setup.py --dataset levir_cd  # prepare list files for a specific dataset
"""

from __future__ import annotations

import argparse
import hashlib
import importlib
import os
import subprocess
import sys
import urllib.request
from pathlib import Path


ROOT = Path(__file__).resolve().parent
DATASET_CONFIGS = ROOT / "configs" / "datasets"
MODEL_REPOS = ROOT / "model_repos"
IMG_EXTS = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp"}
ROOT_LEVEL_REPOS = {
    "BIT_CD",
    "ChangeFormer",
    "STANet",
    "Change3D",
    "BiFA",
    "CDMamba",
    "RSM-CD",
    "SChanger",
    "IFNet",
    "Siam-NestedUNet",
}

REQUIRED_PACKAGES = [
    ("torch", "torch", "all models", False),
    ("torchvision", "torchvision", "all models", False),
    ("numpy", "numpy", "all models", True),
    ("PIL", "Pillow", "all models", True),
    ("cv2", "opencv-python", "all models", True),
    ("yaml", "pyyaml", "all models", True),
    ("tqdm", "tqdm", "weight downloads", True),
    ("timm", "timm", "TinyCD, ELGC-Net, ChangeFormer", True),
    ("einops", "einops", "BIT_CD, ChangeFormer", True),
    ("sklearn", "scikit-learn", "metrics", True),
    ("gdown", "gdown", "GDrive weight fallback", True),
    ("mmengine", "mmengine==0.10.1", "ChangeMamba, Changer", False),
    ("mmcv", "mmcv==2.1.0", "ChangeMamba, Changer", False),
    ("mmseg", "mmsegmentation==1.2.2", "ChangeMamba, Changer", False),
]

CONFIRMED_REPOS = {
    "BIT_CD": "https://github.com/justchenhao/BIT_CD.git",
    "ChangeFormer": "https://github.com/wgcban/ChangeFormer.git",
    "STANet": "https://github.com/justchenhao/STANet.git",
    "Change3D": "https://github.com/Z-Zheng/Change3D.git",
    "BiFA": "https://github.com/zmoka-zht/BiFA.git",
    "CDMamba": "https://github.com/zmoka-zht/CDMamba.git",
    "RSM-CD": "https://github.com/walking-shadow/Official_Remote_Sensing_Mamba.git",
    "SChanger": None,
    "IFNet": "https://github.com/GeoZcx/A-deeply-supervised-image-fusion-network-for-change-detection.git",
    "Siam-NestedUNet": "https://github.com/likyoo/Siam-NestedUNet.git",
    "fully_convolutional_change_detection": "https://github.com/rcdaudt/fully_convolutional_change_detection.git",
    "ChangeMamba": "https://github.com/ChenHongruixuan/ChangeMamba.git",
    "DSAMNet": "https://github.com/liumency/DSAMNet.git",
    "Tiny_model_4_CD": "https://github.com/AndreaCodegoni/Tiny_model_4_CD.git",
    "HANet-CD": "https://github.com/ChengxiHAN/HANet-CD.git",
    "CGNet-CD": "https://github.com/ChengxiHAN/CGNet-CD.git",
    "open-cd": "https://github.com/likyoo/open-cd.git",
    "elgcnet": "https://github.com/techmn/elgcnet.git",
}

ZENODO_WEIGHTS = {
    "vmamba_tiny": {
        "record_id": "14037770",
        "filename": "vssmtiny_dp01_ckpt_epoch_292.pth",
        "dest_dir": MODEL_REPOS / "ChangeMamba" / "pretrained_weight",
        "sha256": None,
        "notes": "VMamba-Tiny for ChangeMamba",
    },
    "vmamba_small": {
        "record_id": "14037770",
        "filename": "vssmsmall_dp03_ckpt_epoch_238.pth",
        "dest_dir": MODEL_REPOS / "ChangeMamba" / "pretrained_weight",
        "sha256": None,
        "notes": "VMamba-Small for ChangeMamba",
    },
    "vmamba_base": {
        "record_id": "14037770",
        "filename": "vssmbase_dp06_ckpt_epoch_241.pth",
        "dest_dir": MODEL_REPOS / "ChangeMamba" / "pretrained_weight",
        "sha256": None,
        "notes": "VMamba-Base for ChangeMamba",
    },
}

TIMM_WEIGHTS_TO_PREFETCH = [
    "efficientnet_b4",
    "mit_b0",
    "mit_b1",
    "mit_b4",
]

TORCHVISION_WEIGHTS_TO_PREFETCH = [
    ("resnet18", "IMAGENET1K_V1"),
    ("resnet50", "IMAGENET1K_V1"),
    ("vgg16", "IMAGENET1K_V1"),
]

TIMM_ALIASES = {
    "efficientnet_b4": ("efficientnet_b4", "tf_efficientnet_b4", "tf_efficientnet_b4_ns"),
    "mit_b0": ("mit_b0", "segformer_b0"),
    "mit_b1": ("mit_b1", "segformer_b1"),
    "mit_b4": ("mit_b4", "segformer_b4"),
}


def run(cmd: list[str]) -> subprocess.CompletedProcess:
    return subprocess.run(cmd, capture_output=True, text=True, check=False)


def import_status(import_name: str) -> tuple[bool, str]:
    try:
        module = importlib.import_module(import_name)
    except Exception as exc:
        return False, str(exc)
    return True, getattr(module, "__version__", "present")


def pip_install(package: str) -> bool:
    print(f"  [INSTALL] pip install {package}")
    result = run([sys.executable, "-m", "pip", "install", package])
    if result.returncode != 0:
        print(f"  [FAIL] {package}: {result.stderr[-500:]}")
        return False
    print(f"  [OK] {package} installed")
    return True


def check_environment(status_only: bool = False) -> dict[str, int]:
    print("\nPYTHON ENVIRONMENT")
    print("------------------")
    print(f"Python executable: {sys.executable}")
    print(f"Python version: {sys.version.split()[0]}")
    present = missing = installed = warnings = 0

    for import_name, pip_name, required_for, auto_install in REQUIRED_PACKAGES:
        ok, detail = import_status(import_name)
        if ok:
            present += 1
            print(f"  [OK] {import_name:<12} {detail:<18} required for {required_for}")
            continue
        missing += 1
        if auto_install and not status_only:
            if pip_install(pip_name):
                installed += 1
                continue
        if not auto_install:
            warnings += 1
            print(f"  [WARN] {import_name:<12} missing; install manually: pip install {pip_name}")
        else:
            print(f"  [MISSING] {import_name:<12} install: pip install {pip_name}")

    if warnings:
        print("\nVersion-sensitive packages are intentionally not auto-installed.")
        print("For ChangeMamba / Changer:")
        print("  pip install mmengine==0.10.1")
        print("  pip install mmcv==2.1.0 -f https://download.openmmlab.com/mmcv/dist/cu124/torch2.6/index.html")
        print("  pip install mmsegmentation==1.2.2 mmdet==3.3.0 mmpretrain==1.2.0")

    return {"present": present, "missing": missing, "installed": installed, "warnings": warnings}


def clone_repo(name: str, url: str | None, dest_dir: Path = MODEL_REPOS, status_only: bool = False) -> bool:
    dest = (ROOT / name) if name in ROOT_LEVEL_REPOS else (dest_dir / name)
    if dest.is_dir():
        print(f"  [OK] {name} already cloned")
        return True
    if url is None:
        print(f"  [SKIP] {name} - URL not confirmed, clone manually")
        return False
    if status_only:
        print(f"  [MISSING] {name} -> {url}")
        return False
    print(f"  [CLONE] {name}...")
    dest.parent.mkdir(parents=True, exist_ok=True)
    result = run(["git", "clone", "--depth", "1", url, str(dest)])
    if result.returncode == 0:
        print(f"  [OK] {name} cloned")
        return True
    print(f"  [FAIL] {name}: {result.stderr[:300]}")
    return False


def clone_repositories(status_only: bool = False) -> dict[str, int]:
    print("\nMODEL REPOSITORIES")
    print("------------------")
    ok = missing = skipped = 0
    for name, url in CONFIRMED_REPOS.items():
        result = clone_repo(name, url, status_only=status_only)
        if result:
            ok += 1
        elif url is None:
            skipped += 1
        else:
            missing += 1
    return {"ok": ok, "missing": missing, "skipped": skipped, "total": len(CONFIRMED_REPOS)}


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


def download_zenodo_weight(
    record_id: str,
    filename: str,
    dest_dir: Path,
    expected_sha256: str | None = None,
    status_only: bool = False,
) -> Path | None:
    dest_dir.mkdir(parents=True, exist_ok=True)
    dest_path = dest_dir / filename
    if dest_path.exists():
        if expected_sha256:
            actual = sha256_file(dest_path)
            if actual != expected_sha256:
                print(f"  [FAIL] {filename} sha256 mismatch")
                return None
        print(f"  [OK] {filename} already present")
        return dest_path
    if status_only:
        print(f"  [MISSING] {filename}")
        return None

    try:
        from tqdm import tqdm
    except Exception:
        tqdm = None

    url = f"https://zenodo.org/records/{record_id}/files/{filename}"
    print(f"  [DOWNLOAD] {filename} from Zenodo record {record_id}")

    class ProgressHook:
        def __init__(self) -> None:
            self.pbar = None

        def __call__(self, count: int, block_size: int, total_size: int) -> None:
            if tqdm is None:
                return
            if self.pbar is None:
                self.pbar = tqdm(total=total_size, unit="B", unit_scale=True, desc=filename)
            self.pbar.update(count * block_size - self.pbar.n)

    hook = ProgressHook()
    try:
        urllib.request.urlretrieve(url, dest_path, reporthook=hook)
        if hook.pbar:
            hook.pbar.close()
        if expected_sha256:
            actual = sha256_file(dest_path)
            if actual != expected_sha256:
                dest_path.unlink(missing_ok=True)
                print(f"  [FAIL] {filename}: sha256 mismatch")
                return None
        print(f"  [OK] {filename} downloaded")
        return dest_path
    except Exception as exc:
        if hook.pbar:
            hook.pbar.close()
        dest_path.unlink(missing_ok=True)
        print(f"  [FAIL] {filename}: {exc}")
        return None


def prefetch_timm_weights(model_name: str, status_only: bool = False) -> bool:
    ok, detail = import_status("timm")
    if not ok:
        print(f"  [SKIP] timm:{model_name} - timm missing ({detail})")
        return False
    if status_only:
        print(f"  [CHECK] timm:{model_name} cache status cannot be proven without instantiation")
        return True

    import timm

    candidates = TIMM_ALIASES.get(model_name, (model_name,))
    for candidate in candidates:
        try:
            model = timm.create_model(candidate, pretrained=True, num_classes=0)
            del model
            print(f"  [OK] timm {model_name} via {candidate}")
            return True
        except RuntimeError as exc:
            if "Unknown model" not in str(exc):
                print(f"  [FAIL] timm {model_name}: {exc}")
                return False
    print(f"  [WARN] timm {getattr(timm, '__version__', 'unknown')} does not provide {model_name}")
    return False


def prefetch_torchvision_weights(model_name: str, weight_name: str, status_only: bool = False) -> bool:
    ok, detail = import_status("torchvision.models")
    if not ok:
        print(f"  [SKIP] torchvision:{model_name} - torchvision missing ({detail})")
        return False
    if status_only:
        print(f"  [CHECK] torchvision:{model_name} cache status cannot be proven without instantiation")
        return True

    import torchvision.models as models

    builder = getattr(models, model_name)
    enum_name = "".join(part.capitalize() for part in model_name.split("_")) + "_Weights"
    weights_enum = getattr(models, enum_name, None)
    try:
        if weights_enum is not None:
            weights = getattr(weights_enum, weight_name, weights_enum.DEFAULT)
            model = builder(weights=weights)
        else:
            model = builder(pretrained=True)
        del model
        print(f"  [OK] torchvision {model_name}")
        return True
    except Exception as exc:
        print(f"  [FAIL] torchvision {model_name}: {exc}")
        return False


def download_weights(status_only: bool = False) -> dict[str, int]:
    print("\nPRETRAINED WEIGHTS")
    print("------------------")
    ok = missing = 0
    for spec in ZENODO_WEIGHTS.values():
        path = download_zenodo_weight(
            spec["record_id"],
            spec["filename"],
            spec["dest_dir"],
            expected_sha256=spec["sha256"],
            status_only=status_only,
        )
        ok += int(path is not None)
        missing += int(path is None)
    for model_name in TIMM_WEIGHTS_TO_PREFETCH:
        ok += int(prefetch_timm_weights(model_name, status_only=status_only))
    for model_name, weight_name in TORCHVISION_WEIGHTS_TO_PREFETCH:
        ok += int(prefetch_torchvision_weights(model_name, weight_name, status_only=status_only))
    return {"ok": ok, "missing": missing, "total": len(ZENODO_WEIGHTS) + len(TIMM_WEIGHTS_TO_PREFETCH) + len(TORCHVISION_WEIGHTS_TO_PREFETCH)}


def parse_simple_yaml(path: Path) -> dict[str, object]:
    data: dict[str, object] = {}
    stack: list[tuple[int, dict[str, object]]] = [(-1, data)]
    for raw in path.read_text(encoding="utf-8").splitlines():
        line = raw.split("#", 1)[0].rstrip()
        if not line.strip() or ":" not in line:
            continue
        indent = len(line) - len(line.lstrip(" "))
        key, value = line.strip().split(":", 1)
        while stack and indent <= stack[-1][0]:
            stack.pop()
        current = stack[-1][1]
        value = value.strip()
        if not value:
            child: dict[str, object] = {}
            current[key] = child
            stack.append((indent, child))
            continue
        if value.startswith("[") and value.endswith("]"):
            current[key] = [x.strip().strip("'\"") for x in value[1:-1].split(",") if x.strip()]
        elif value.isdigit():
            current[key] = int(value)
        else:
            current[key] = value.strip("'\"")
    return data


def dataset_root_from_config(cfg: dict[str, object]) -> Path:
    value = str(cfg.get("data_root", ""))
    data_root = os.environ.get("DATA_ROOT", str(ROOT.parent / "Datasets"))
    return Path(value.replace("${DATA_ROOT}", data_root)).expanduser()


def prepare_dataset_lists(data_root: Path, image_a_folder: str = "A", splits: tuple[str, ...] = ("train", "val", "test")) -> dict[str, int]:
    list_dir = data_root / "list"
    list_dir.mkdir(parents=True, exist_ok=True)
    counts: dict[str, int] = {}
    for split in splits:
        out = list_dir / f"{split}.txt"
        if out.exists():
            counts[split] = sum(1 for _ in out.open("r", encoding="utf-8"))
            print(f"  [OK] {out} already exists ({counts[split]} entries)")
            continue
        split_dir = data_root / split
        img_dir = None
        for candidate in (image_a_folder, "A", "T1", "t1", "img", "images", "image"):
            if (split_dir / candidate).is_dir():
                img_dir = split_dir / candidate
                break
        if img_dir is None:
            counts[split] = 0
            print(f"  [WARN] no image folder found under {split_dir}")
            continue
        names = sorted(p.name for p in img_dir.iterdir() if p.is_file() and p.suffix.lower() in IMG_EXTS)
        out.write_text("\n".join(names) + ("\n" if names else ""), encoding="utf-8")
        counts[split] = len(names)
        print(f"  [OK] {out}: {len(names)} entries")
    return counts


def prepare_datasets(selected: str | None = None, status_only: bool = False) -> dict[str, int]:
    print("\nDATASETS")
    print("--------")
    exists = missing = prepared = 0
    paths = sorted(DATASET_CONFIGS.glob("*.yaml"))
    for path in paths:
        if selected and path.stem != selected:
            continue
        cfg = parse_simple_yaml(path)
        root = dataset_root_from_config(cfg)
        if root.is_dir():
            exists += 1
            print(f"  [EXISTS] {path.stem}: {root}")
            if not status_only:
                split_keys = tuple((cfg.get("splits") or {"train": "train", "val": "val", "test": "test"}).keys())
                prepare_dataset_lists(root, image_a_folder=str(cfg.get("image_a_folder", "A")), splits=split_keys)
                prepared += 1
        else:
            missing += 1
            print(f"  [MISSING] {path.stem}: {root}")
    return {"exists": exists, "missing": missing, "prepared": prepared, "total": len(paths)}


def print_final_status(env: dict[str, int], repos: dict[str, int], weights: dict[str, int] | None, datasets: dict[str, int]) -> None:
    print("\nSETUP STATUS")
    print("Component                     Status        Notes")
    print("-----------------------------|-------------|----------------------------------")
    env_status = "COMPLETE" if env["missing"] == 0 else "CHECK"
    print(f"Python environment            {env_status:<12} {env['installed']} packages auto-installed; {env['warnings']} manual warnings")
    print(f"Model repos ({repos['total']} total)       {repos['ok']} / {repos['total']:<6} {repos['skipped']} URL-unconfirmed entries")
    if weights is None:
        print("Backbone weights              SKIPPED       run: python setup.py --weights-only")
    else:
        print(f"Backbone weights              {weights['ok']} / {weights['total']:<6} Zenodo, timm, and torchvision warmups")
    print(f"Datasets                      {datasets['exists']} / {datasets['total']:<6} roots found under DATA_ROOT/default search")
    print("\nNext steps:")
    print("  export DATA_ROOT=/path/to/datasets")
    print("  python run_training.py --model bifa --dataset levir_cd --dry-run")
    print("  python run_training.py --model bifa --dataset levir_cd")
    print("  For ChangeMamba/Changer, install the MMSeg stack shown above.")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="One-command setup for the CD-Models benchmark suite.")
    parser.add_argument("--skip-weights", action="store_true")
    parser.add_argument("--weights-only", action="store_true")
    parser.add_argument("--env-check-only", action="store_true")
    parser.add_argument("--status", action="store_true")
    parser.add_argument("--dataset", default=None, help="Prepare list files for one dataset config name.")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    status_only = args.status

    env = check_environment(status_only=status_only)
    if args.env_check_only:
        print_final_status(env, {"ok": 0, "missing": 0, "skipped": 0, "total": len(CONFIRMED_REPOS)}, None, {"exists": 0, "missing": 0, "prepared": 0, "total": 0})
        return 0 if env["warnings"] == 0 else 1

    repos = {"ok": 0, "missing": 0, "skipped": 0, "total": len(CONFIRMED_REPOS)}
    if not args.weights_only:
        repos = clone_repositories(status_only=status_only)

    weights = None
    if not args.skip_weights:
        weights = download_weights(status_only=status_only)

    datasets = prepare_datasets(selected=args.dataset, status_only=status_only)
    print_final_status(env, repos, weights, datasets)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())