#!/usr/bin/env python3
import argparse
from pathlib import Path


def ensure_link(link: Path, target: str):
    if link.exists() or link.is_symlink():
        return
    try:
        link.symlink_to(target)
    except OSError:
        print(f"Could not create symlink {link} -> {target}. Please create it manually.")


def main():
    parser = argparse.ArgumentParser(description="Prepare DRSeg compatibility links.")
    parser.add_argument("--data-root", default="data/DRSeg")
    args = parser.parse_args()

    root = Path(args.data_root)
    if not root.exists():
        raise SystemExit(f"Missing data root: {root}")

    ensure_link(root / "CODrone", ".")
    if (root / "label").exists():
        ensure_link(root / "labels", "label")

    required = [
        root / "DRtrain",
        root / "DRval",
        root / "DRtest",
        root / "label" / "DRSeg_train.json",
        root / "label" / "DRSeg_val.json",
        root / "label" / "DRSeg_test.json",
    ]
    missing = [str(p) for p in required if not p.exists()]
    if missing:
        raise SystemExit("Missing required DRSeg files:\n" + "\n".join(missing))
    print(f"DRSeg is ready: {root}")


if __name__ == "__main__":
    main()