import argparse

import torch
from sentence_transformers import SentenceTransformer


def detect_devices():
    devices = []

    if torch.cuda.is_available():
        devices.extend(f"cuda:{index}" for index in range(torch.cuda.device_count()))

    mps_backend = getattr(torch.backends, "mps", None)
    if mps_backend is not None and mps_backend.is_available():
        devices.append("mps")

    devices.append("cpu")
    return devices


parser = argparse.ArgumentParser()
parser.add_argument("--device", default="auto", help="auto, cpu, cuda:0, or mps")
args = parser.parse_args()

available_devices = detect_devices()
requested_device = args.device

if requested_device == "auto":
    selected_device = available_devices[0]
elif requested_device in available_devices:
    selected_device = requested_device
else:
    raise SystemExit(
        f"{requested_device} is not available. Available devices: "
        f"{', '.join(available_devices)}"
    )

model = SentenceTransformer(
    "sentence-transformers/all-MiniLM-L6-v2",
    device=selected_device,
)

embeddings = model.encode(
    [
        "Route embedding inference to the selected accelerator.",
        "Keep a CPU fallback for development machines.",
    ],
    show_progress_bar=False,
)

print(f"requested device: {requested_device}")
print(f"available devices: {', '.join(available_devices)}")
print(f"selected device: {model.device}")
print(f"embedding shape: {embeddings.shape}")