from __future__ import annotations

from pathlib import Path

import joblib
from sentence_transformers import SentenceTransformer
from sklearn.neighbors import KNeighborsClassifier

EMBEDDER_MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2"


def predict_intent(
    embedder: SentenceTransformer, classifier: KNeighborsClassifier, text: str
) -> str:
    embedding = embedder.encode([text])
    prediction = classifier.predict(embedding)[0]
    return str(prediction)


def load_artifacts(
    artifacts_dir: Path,
) -> tuple[KNeighborsClassifier, dict[str, float]]:
    classifier = joblib.load(artifacts_dir / "classifier.joblib")
    return classifier


def main() -> None:
    base_dir = Path(__file__).resolve().parent.parent
    artifacts_dir = base_dir / "artifacts"

    classifier = load_artifacts(artifacts_dir)
    embedder = SentenceTransformer(EMBEDDER_MODEL_ID)

    sample = "book a flight for next friday"
    predicted_intent = predict_intent(embedder, classifier, sample)

    print(f"Sample: {sample}")
    print(f"Predicted intent: {predicted_intent}")


if __name__ == "__main__":
    main()