"""Fine-tune the email-risk classifier (HF Pro).

Trains distilbert-base-uncased (4 classes) on the blocklist-derived dataset.
The dataset is class-imbalanced (DISPOSABLE >> SPAM_TRAP), so cross-entropy is
class-weighted and per-class precision is reported - the promotion gate.

Usage:
    python -m app.ml.train --data app/ml/data/email_risk.csv \
        --hub-repo <hf-user>/email-risk-distilbert [--eval-only --model <repo>]

Requires requirements-ml.txt (torch CPU, transformers, datasets, evaluate).
"""
from __future__ import annotations

import argparse
import json
import sys
from collections import Counter
from pathlib import Path

LABEL_NAMES = ["NORMAL", "DISPOSABLE", "ROLE", "SPAM_TRAP"]
PROMOTION_PRECISION = 0.98  # DISPOSABLE precision required to leave advisory


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data", default="app/ml/data/email_risk.csv")
    parser.add_argument("--hub-repo", default="")
    parser.add_argument("--model", default="distilbert-base-uncased")
    parser.add_argument("--eval-only", action="store_true")
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--out", default="app/ml/metrics.json")
    args = parser.parse_args()

    import datasets
    import evaluate
    import numpy as np
    import torch
    from sklearn.metrics import precision_recall_fscore_support
    from transformers import (
        AutoModelForSequenceClassification,
        AutoTokenizer,
        Trainer,
        TrainingArguments,
    )

    raw = datasets.load_dataset("csv", data_files=args.data)["train"]
    raw = raw.train_test_split(test_size=0.15, seed=42)

    tokenizer = AutoTokenizer.from_pretrained(args.model)

    def tok(batch):
        return tokenizer(batch["text"], padding="max_length", truncation=True,
                         max_length=64)

    ds = raw.map(tok, batched=True)
    labels = ds["train"]["label"]
    counts = Counter(labels)
    total = sum(counts.values())
    weights = [total / counts.get(i, 1) for i in range(len(LABEL_NAMES))]
    print(f"class counts: {dict(sorted(counts.items()))} -> weights {weights}")

    model = AutoModelForSequenceClassification.from_pretrained(
        args.model, num_labels=len(LABEL_NAMES),
        id2label={i: n for i, n in enumerate(LABEL_NAMES)},
        label2id={n: i for i, n in enumerate(LABEL_NAMES)},
    )

    def compute_metrics(eval_pred):
        logits, y = eval_pred
        pred = np.argmax(logits, axis=-1)
        p, r, f1, _ = precision_recall_fscore_support(
            y, pred, labels=range(len(LABEL_NAMES)), zero_division=0)
        return {
            "accuracy": float((pred == y).mean()),
            **{f"precision_{LABEL_NAMES[i]}": float(p[i]) for i in range(len(LABEL_NAMES))},
            **{f"recall_{LABEL_NAMES[i]}": float(r[i]) for i in range(len(LABEL_NAMES))},
        }

    class WeightedTrainer(Trainer):
        def compute_loss(self, model, inputs, return_outputs=False, **kw):
            y = inputs.pop("labels")
            outputs = model(**inputs)
            loss = torch.nn.functional.cross_entropy(
                outputs.logits, y, weight=torch.tensor(weights, dtype=outputs.logits.dtype,
                                                       device=outputs.logits.device))
            return (loss, outputs) if return_outputs else loss

    targs = TrainingArguments(
        output_dir="app/ml/runs",
        learning_rate=2e-5,
        per_device_train_batch_size=64,
        per_device_eval_batch_size=128,
        num_train_epochs=args.epochs,
        weight_decay=0.01,
        eval_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
        use_cpu=not torch.cuda.is_available(),
        hub_model_id=args.hub_repo or None,
        hub_strategy="end",
        push_to_hub=bool(args.hub_repo),
    )

    trainer = WeightedTrainer(model=model, args=targs, train_dataset=ds["train"],
                              eval_dataset=ds["test"], compute_metrics=compute_metrics)

    if args.eval_only:
        metrics = trainer.evaluate()
    else:
        trainer.train()
        metrics = trainer.evaluate()
        if args.hub_repo:
            trainer.push_to_hub()

    disposable_precision = metrics.get("eval_precision_DISPOSABLE", 0.0)
    result = {
        "model": args.hub_repo or args.model,
        "metrics": {k: round(v, 4) for k, v in metrics.items() if k.startswith("eval_")},
        "class_weights": weights,
        "promotion": {
            "gate": f"precision_DISPOSABLE >= {PROMOTION_PRECISION}",
            "precision_disposable": round(disposable_precision, 4),
            "promotion_ready": disposable_precision >= PROMOTION_PRECISION,
        },
    }
    Path(args.out).write_text(json.dumps(result, indent=2), encoding="utf-8")
    print(json.dumps(result["promotion"], indent=2))
    print(f"metrics -> {args.out}")
    print("PROMOTION READY" if result["promotion"]["promotion_ready"] else
          "STAY ADVISORY - precision gate not met")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())