from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Dict, List

import numpy as np
from sentence_transformers import SentenceTransformer


ROOT = Path(__file__).resolve().parent


class AgentWeaveSemanticRouter:
    """Prototype-based semantic capability router built on frozen MiniLM embeddings.

    This is an experimental semantic companion to AgentWeave's default
    deterministic routing path. It does not replace AgentWeave policy,
    authorization, or execution controls. The upstream encoder is loaded as a
    runtime dependency; this repository is not a fine-tuned MiniLM model.
    """

    def __init__(
        self,
        config_path: str | Path = ROOT / "config.json",
        prototypes_path: str | Path = ROOT / "route_prototypes.json",
    ) -> None:
        self.config = json.loads(Path(config_path).read_text(encoding="utf-8"))
        self.prototypes: Dict[str, List[str]] = json.loads(
            Path(prototypes_path).read_text(encoding="utf-8")
        )
        encoder_model = self.config.get("encoder_model") or self.config.get("base_model")
        if not encoder_model:
            raise ValueError("config.json must define 'encoder_model'")
        self.model = SentenceTransformer(str(encoder_model), device="cpu")

        texts: List[str] = []
        labels: List[str] = []
        for label, examples in self.prototypes.items():
            for example in examples:
                labels.append(label)
                texts.append(example)

        self._prototype_labels = labels
        self._prototype_embeddings = self.model.encode(
            texts,
            normalize_embeddings=bool(self.config.get("normalize_embeddings", True)),
            convert_to_numpy=True,
            show_progress_bar=False,
        )

    def route(self, query: str, top_k: int | None = None) -> List[dict]:
        if not query or not query.strip():
            raise ValueError("query must be a non-empty string")

        top_k = int(top_k or self.config.get("default_top_k", 3))
        query_embedding = self.model.encode(
            [query],
            normalize_embeddings=bool(self.config.get("normalize_embeddings", True)),
            convert_to_numpy=True,
            show_progress_bar=False,
        )[0]

        similarities = self._prototype_embeddings @ query_embedding
        best_by_label: Dict[str, float] = {}
        for label, score in zip(self._prototype_labels, similarities):
            best_by_label[label] = max(best_by_label.get(label, -1.0), float(score))

        ranked = sorted(best_by_label.items(), key=lambda item: item[1], reverse=True)
        return [
            {"route": label, "score": round(score, 6)}
            for label, score in ranked[: max(1, min(top_k, len(ranked)))]
        ]


def main() -> None:
    parser = argparse.ArgumentParser(description="AgentWeave MiniLM semantic router")
    parser.add_argument("query", help="Task or request to route")
    parser.add_argument("--top-k", type=int, default=None, help="Number of routes to return")
    args = parser.parse_args()

    router = AgentWeaveSemanticRouter()
    print(json.dumps(router.route(args.query, args.top_k), indent=2))


if __name__ == "__main__":
    main()