"""Probe the self-hosted LLM endpoint (aipi).

Lists available models and verifies the embedding dimension of the
configured ``embed`` model against ``BOR_EMBEDDING_DIM`` (default 768).
Run this before the first import if the LLM backend ever changes:

    uv run python -m scripts.llm_probe
"""
from __future__ import annotations

import os
import sys

import httpx
from dotenv import load_dotenv

load_dotenv()


def main() -> int:
    base_url = os.environ.get("BOR_LLM_BASE_URL", "https://aipi.reeseapps.com/v1").rstrip("/")
    api_key = (
        os.environ.get("BOR_LLM_API_KEY")
        or os.environ.get("AIPI_KEY")
        or "not-needed"
    )
    embed_model = os.environ.get("BOR_LLM_EMBED_MODEL", "embed")
    chat_model = os.environ.get("BOR_LLM_CHAT_MODEL", "turbo")
    expected_dim = int(os.environ.get("BOR_EMBEDDING_DIM", "768"))

    headers = {"Authorization": f"Bearer {api_key}"}
    with httpx.Client(base_url=base_url, headers=headers, timeout=30.0) as client:
        r = client.get("/models")
        r.raise_for_status()
        models = [m["id"] for m in r.json()["data"]]
        print(f"[probe] base_url : {base_url}")
        print(f"[probe] models   : {', '.join(models)}")

        for needed in (chat_model, embed_model):
            if needed not in models:
                print(f"[probe] ERROR: required model '{needed}' not available")
                return 1

        r = client.post(
            "/embeddings",
            json={"model": embed_model, "input": "brain of reese dimension probe"},
        )
        r.raise_for_status()
        dims = sorted({len(d["embedding"]) for d in r.json()["data"]})
        print(f"[probe] dims({embed_model}): {dims}")

    if dims != [expected_dim]:
        print(
            f"[probe] MISMATCH: expected {expected_dim}, got {dims}. "
            "Update BOR_EMBEDDING_DIM and recreate the chunks table (see README)."
        )
        return 1
    print("[probe] OK — models present, embedding dimension matches configuration.")
    return 0


if __name__ == "__main__":
    sys.exit(main())