"""Custom handler for Hugging Face Inference Endpoints.

HF creates EndpointHandler(path) once when the endpoint starts, then calls it
for every request with the JSON body, e.g. {"inputs": "hello how are you"}.
"""
import os
from typing import Any, Dict, List, Union

import joblib
from sentence_transformers import SentenceTransformer

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

# What to say for each label the classifier can predict.
# Every label used in data.csv needs an entry here.
REPLIES = {
    "greeting": "hello",
    "computer": "output",
    "other": "I don't understand",
}


class EndpointHandler:
    def __init__(self, path: str = ""):
        base = path or os.path.dirname(os.path.abspath(__file__))
        self.encoder = SentenceTransformer(ENCODER_NAME)
        self.clf = joblib.load(os.path.join(base, "classifier.joblib"))

    def _predict(self, texts: List[str]) -> List[Dict[str, Any]]:
        if not texts:
            return []
        embeddings = self.encoder.encode(texts, normalize_embeddings=True)
        probs = self.clf.predict_proba(embeddings)
        results = []
        for row in probs:
            best = int(row.argmax())
            label = str(self.clf.classes_[best])
            results.append({
                "reply": REPLIES[label],
                "label": label,
                "score": round(float(row[best]), 4),
            })
        return results

    def __call__(self, data: Dict[str, Any]) -> Union[Dict[str, Any], List[Dict[str, Any]]]:
        if not isinstance(data, dict) or "inputs" not in data:
            raise ValueError('Send JSON like {"inputs": "hi"}')
        inputs = data["inputs"]
        if isinstance(inputs, list):  # several texts -> list of results
            return self._predict([str(t) for t in inputs])
        return self._predict([str(inputs)])[0]  # one text -> one result