"""
Train a tiny text classifier on the generated emails and save it to disk.

A model here is a saved recipe: it learned which words tend to go with
which labels. Training means showing it labeled examples. Predict means
asking it to guess a label for new text. It does not “understand” email
the way a person does, and it never calls an LLM or the internet.
"""

from __future__ import annotations

import json
from pathlib import Path

import joblib
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline

HERE = Path(__file__).resolve().parent
TRAINING_PATH = HERE / "data" / "training.json"
MODEL_DIR = HERE / "model"
MODEL_PATH = MODEL_DIR / "inbox_model.joblib"


def load_training(path: Path = TRAINING_PATH) -> tuple[list[str], list[str]]:
    if not path.exists():
        raise FileNotFoundError(
            f"Missing {path}. Run: python generate_data.py"
        )
    rows = json.loads(path.read_text(encoding="utf-8"))
    texts = []
    labels = []
    for row in rows:
        # Combine subject + body so the model sees the whole message.
        texts.append(f"{row['subject']} {row['body']}")
        labels.append(row["label"])
    return texts, labels


def build_pipeline() -> Pipeline:
    """
    Pipeline = two steps glued together.

    1. TfidfVectorizer turns words into numbers (how important a word is
       in this message compared with the rest of the training set).
    2. LogisticRegression learns a weight for each word per label, then
       picks the label with the strongest score.
    """
    return Pipeline(
        [
            ("tfidf", TfidfVectorizer()),
            ("clf", LogisticRegression(max_iter=1000)),
        ]
    )


def train_and_save(
    training_path: Path = TRAINING_PATH,
    model_path: Path = MODEL_PATH,
) -> tuple[Pipeline, float, int]:
    texts, labels = load_training(training_path)

    # Hold out some labeled rows so we can check the recipe on emails
    # it did not use while learning. random_state keeps the split the same
    # every run (deterministic, like generate_data.py).
    X_train, X_valid, y_train, y_valid = train_test_split(
        texts,
        labels,
        test_size=0.25,
        random_state=42,
        stratify=labels,
    )

    pipeline = build_pipeline()
    # fit = train: look at X_train and y_train and store patterns.
    pipeline.fit(X_train, y_train)
    # predict = guess labels for the held-out messages.
    guessed = pipeline.predict(X_valid)
    accuracy = float(accuracy_score(y_valid, guessed))

    model_path.parent.mkdir(parents=True, exist_ok=True)
    joblib.dump(pipeline, model_path)
    return pipeline, accuracy, len(y_valid)


def main() -> None:
    _pipeline, accuracy, n_valid = train_and_save()
    print(f"Saved model to {MODEL_PATH}")
    print(f"Validation accuracy: {accuracy:.0%} ({n_valid} held-out generated emails)")
    print(
        "Read this number carefully: it only describes this practice dataset. "
        "It is not a claim that the model works on real customer inboxes."
    )
    print("Done. Next: python agent.py")


if __name__ == "__main__":
    main()