from __future__ import annotations

import argparse
import json
from pathlib import Path

from sentence_transformers import SentenceTransformer


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate semantic embeddings")
    parser.add_argument(
        "--registry-path",
        default="src/rabbitmq_mcp/data/operations.json",
    )
    parser.add_argument(
        "--output-path",
        default="src/rabbitmq_mcp/data/embeddings.json",
    )
    parser.add_argument(
        "--model-name",
        default="all-MiniLM-L6-v2",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    registry_path = Path(args.registry_path)
    registry = json.loads(registry_path.read_text(encoding="utf-8"))
    operations = registry.get("operations", {})
    operation_ids = list(operations.keys())
    texts = []
    for op_id in operation_ids:
        op = operations[op_id]
        summary = op.get("summary") or ""
        description = op.get("description") or ""
        texts.append(f"{op_id} {summary} {description}".strip())

    model = SentenceTransformer(args.model_name)
    vectors = model.encode(texts, normalize_embeddings=True)
    embeddings = {
        op_id: vectors[idx].astype(float).tolist()
        for idx, op_id in enumerate(operation_ids)
    }
    payload = {
        "model": args.model_name,
        "dimension": int(vectors.shape[1]),
        "embeddings": embeddings,
    }
    output_path = Path(args.output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")


if __name__ == "__main__":
    main()
