"""Verify SamGeo3 / PyTorch / CUDA install for this isolated lab env.

Usage (from project root, with venv active):
    python scripts/check_install.py
    python scripts/check_install.py --load-model
"""

from __future__ import annotations

import argparse
import os
import sys
import traceback


def _ok(msg: str) -> None:
    print(f"[OK]   {msg}")


def _warn(msg: str) -> None:
    print(f"[WARN] {msg}")


def _fail(msg: str) -> None:
    print(f"[FAIL] {msg}")


def check_torch() -> bool:
    try:
        import torch
    except ImportError as e:
        _fail(f"torch import failed: {e}")
        return False

    _ok(f"torch {torch.__version__}")
    cuda = torch.cuda.is_available()
    if cuda:
        name = torch.cuda.get_device_name(0)
        cap = torch.cuda.get_device_capability(0)
        _ok(f"CUDA available | GPU={name} | capability={cap}")
        _ok(f"cuda runtime reported by torch: {torch.version.cuda}")
    else:
        _warn("CUDA not available — SAM3 meta backend needs NVIDIA GPU + CUDA torch")
    return True


def check_samgeo3_imports() -> bool:
    try:
        import samgeo
        from samgeo import SamGeo3
    except ImportError as e:
        _fail(f"samgeo / SamGeo3 import failed: {e}")
        return False

    _ok(f"samgeo {getattr(samgeo, '__version__', 'unknown')}")
    _ok("SamGeo3 class importable")

    try:
        import sam3  # noqa: F401

        _ok(f"sam3 package present ({getattr(sam3, '__version__', 'no __version__')})")
    except ImportError as e:
        _fail(f"sam3 package missing (required for backend='meta'): {e}")
        return False

    return True


def check_hf_access() -> None:
    token_env = bool(os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN"))
    token_file = os.path.expanduser("~/.cache/huggingface/token")
    has_file = os.path.isfile(token_file)
    if token_env:
        _ok("HF token found in environment")
    elif has_file:
        _ok(f"HF token file present: {token_file}")
    else:
        _warn(
            "No HF token found. SAM3/3.1 gated models need: "
            "hf auth login  (after HF access approval)"
        )

    ckpt = os.environ.get("SAM3_CHECKPOINT_PATH")
    if ckpt:
        exists = os.path.isfile(ckpt)
        ( _ok if exists else _fail)(f"SAM3_CHECKPOINT_PATH={ckpt} exists={exists}")
    else:
        _warn("SAM3_CHECKPOINT_PATH not set (optional; HF download/cache used otherwise)")


def try_load_model(model_id: str, device: str | None) -> bool:
    try:
        from samgeo import SamGeo3
        import torch
    except ImportError as e:
        _fail(f"imports for model load failed: {e}")
        return False

    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"

    if device == "cuda" and not torch.cuda.is_available():
        _fail("requested CUDA but torch.cuda.is_available() is False")
        return False

    print(f"\nLoading SamGeo3(backend='meta', model_id='{model_id}', device='{device}') ...")
    try:
        kwargs = {
            "backend": "meta",
            "model_id": model_id,
            "device": device,
            "confidence_threshold": 0.5,
            "enable_segmentation": True,
            "enable_inst_interactivity": False,
        }
        ckpt = os.environ.get("SAM3_CHECKPOINT_PATH")
        if ckpt and os.path.isfile(ckpt):
            kwargs["checkpoint_path"] = ckpt
            kwargs["load_from_HF"] = False
            _ok(f"using local checkpoint: {ckpt}")

        sam = SamGeo3(**kwargs)
        _ok(f"model loaded | backend={sam.backend} | device={getattr(sam, 'device', device)}")
        del sam
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        return True
    except Exception as e:
        _fail(f"model load failed: {e}")
        traceback.print_exc()
        return False


def main() -> int:
    parser = argparse.ArgumentParser(description="SamGeo3 install verification")
    parser.add_argument(
        "--load-model",
        action="store_true",
        help="Also instantiate SamGeo3 (downloads weights if needed)",
    )
    parser.add_argument(
        "--model-id",
        default="facebook/sam3.1",
        help="HF model id for --load-model (default: facebook/sam3.1)",
    )
    parser.add_argument(
        "--device",
        default=None,
        help="cuda | cpu (default: auto)",
    )
    args = parser.parse_args()

    print("=== SamGeo3 lab: install check ===\n")
    print(f"python: {sys.version}")
    print(f"exe:    {sys.executable}\n")

    ok = True
    ok = check_torch() and ok
    ok = check_samgeo3_imports() and ok
    check_hf_access()

    if args.load_model:
        ok = try_load_model(args.model_id, args.device) and ok

    print()
    if ok:
        _ok("all required checks passed")
        return 0
    _fail("one or more required checks failed")
    return 1


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