from __future__ import annotations

import csv
from pathlib import Path

import joblib
from sentence_transformers import SentenceTransformer
from sklearn.metrics import accuracy_score, classification_report
from sklearn.model_selection import train_test_split
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_dataset(dataset_path: Path) -> tuple[list[str], list[str]]:
    texts: list[str] = []
    labels: list[str] = []

    with dataset_path.open("r", encoding="utf-8", newline="") as csv_file:
        reader = csv.DictReader(csv_file)
        for row in reader:
            text = row.get("text", "").strip()
            label = row.get("label", "").strip()
            if not text or not label:
                continue
            texts.append(text)
            labels.append(label)

    if not texts or not labels:
        raise ValueError(f"Dataset is empty or invalid: {dataset_path}")

    return texts, labels


def train_model(
    texts: list[str], labels: list[str], random_state: int = 42
) -> tuple[SentenceTransformer, KNeighborsClassifier, dict[str, float]]:
    x_train, x_test, y_train, y_test = train_test_split(
        texts,
        labels,
        test_size=0.3,
        random_state=random_state,
        stratify=labels,
    )

    embedder = SentenceTransformer(EMBEDDER_MODEL_ID)
    x_train_emb = embedder.encode(x_train)
    x_test_emb = embedder.encode(x_test)

    classifier = KNeighborsClassifier(n_neighbors=3)
    classifier.fit(x_train_emb, y_train)

    y_pred = classifier.predict(x_test_emb)
    accuracy = float(accuracy_score(y_test, y_pred))

    print("Classification report:")
    print(classification_report(y_test, y_pred))

    return embedder, classifier, {"accuracy": accuracy}


def save_artifacts(
    classifier: KNeighborsClassifier,
    metrics: dict[str, float],
    output_dir: Path,
) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)

    joblib.dump(classifier, output_dir / "classifier.joblib")


def main() -> None:
    base_dir = Path(__file__).resolve().parent.parent
    dataset_path = base_dir / "data" / "intents.csv"
    output_dir = base_dir / "artifacts"

    texts, labels = load_dataset(dataset_path)
    embedder, classifier, metrics = train_model(texts, labels)
    save_artifacts(classifier, metrics, output_dir)

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

    print(f"Accuracy: {metrics['accuracy']:.4f}")
    print(f"Sample: {sample}")
    print(f"Predicted intent: {predicted_intent}")


if __name__ == "__main__":
    main()