#!/usr/bin/env python3
"""Quick CLI smoke test for a self-contained Galileo checkpoint folder."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import torch
from transformers import pipeline


def parse_args() -> argparse.Namespace:
    repo_root = Path(__file__).resolve().parent
    parser = argparse.ArgumentParser(description="Run Galileo feature extraction on dummy or real input.")
    parser.add_argument(
        "--model",
        type=Path,
        default=repo_root / "galileo-nano-patch8",
        help="Path to a checkpoint folder (default: galileo-nano-patch8)",
    )
    parser.add_argument(
        "--size",
        type=int,
        default=32,
        help="Spatial size H=W for dummy S2 input (default: 32)",
    )
    parser.add_argument(
        "--pool",
        action="store_true",
        default=True,
        help="Return pooled features (default: True)",
    )
    parser.add_argument(
        "--no-pool",
        action="store_true",
        help="Return sequence features instead of pooled output",
    )
    parser.add_argument(
        "--device",
        default=None,
        help="Torch device, e.g. cuda or cpu (default: auto)",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    model_dir = args.model.resolve()
    if not model_dir.is_dir():
        raise SystemExit(f"Model folder not found: {model_dir}")

    with open(model_dir / "config.json", encoding="utf-8") as handle:
        config = json.load(handle)

    hidden_size = config["hidden_size"]
    device = args.device or ("cuda" if torch.cuda.is_available() else "cpu")
    pool = not args.no_pool

    # 10-band Sentinel-2 stack (galileo S2_BANDS ordering)
    s2 = np.random.randn(args.size, args.size, 1, 10).astype(np.float32)

    print(f"model:  {model_dir}")
    print(f"input:  dummy S2 array shape={s2.shape}")
    print(f"device: {device}")
    print(f"pool:   {pool}")

    pipe = pipeline(
        task="galileo-feature-extraction",
        model=str(model_dir),
        trust_remote_code=True,
        device=device,
    )

    features = pipe(s2=s2, pool=pool, return_tensors=True, normalize=False)
    print(f"output: {tuple(features.shape)} dtype={features.dtype}")
    assert features.shape[-1] == hidden_size
    print("OK")


if __name__ == "__main__":
    main()