"""Türkçe prompt injection ikili sınıflandırıcısını eğitir (encoder, softmax iki sınıf).

Kullanım:
  python -m pinj.train --base models/berturk-cased --data data/pinj --out runs/pinj-berturk
Çıktı: runs/.../final (model + tokenizer), runs/.../train_report.json (dev metrikleri)
"""
import argparse
import json
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import Dataset
from transformers import (AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments,
                          set_seed)

ID2LABEL = {0: "BENIGN", 1: "INJECTION"}
LABEL2ID = {v: k for k, v in ID2LABEL.items()}


class PinjData(Dataset):
    def __init__(self, path, tok, max_len):
        rows = [json.loads(l) for l in open(path, encoding="utf-8")]
        self.enc = tok([r["text"] for r in rows], truncation=True, max_length=max_len)
        self.y = np.array([r["label"] for r in rows], dtype=np.int64)

    def __len__(self):
        return len(self.y)

    def __getitem__(self, i):
        item = {k: v[i] for k, v in self.enc.items()}
        item["labels"] = self.y[i]
        return item


def collate(batch, pad_id):
    n = max(len(b["input_ids"]) for b in batch)
    out = {}
    for k in batch[0]:
        if k == "labels":
            out[k] = torch.tensor([b[k] for b in batch])
        else:
            pad = pad_id if k == "input_ids" else 0
            out[k] = torch.tensor([list(b[k]) + [pad] * (n - len(b[k])) for b in batch])
    return out


def prf(y, p):
    tp = int(((p == 1) & (y == 1)).sum())
    fp = int(((p == 1) & (y == 0)).sum())
    fn = int(((p == 0) & (y == 1)).sum())
    pr = tp / (tp + fp) if tp + fp else 0.0
    rc = tp / (tp + fn) if tp + fn else 0.0
    return pr, rc, (2 * pr * rc / (pr + rc) if pr + rc else 0.0)


def compute_metrics(pred):
    p = pred.predictions.argmax(-1)
    pr, rc, f1 = prf(pred.label_ids, p)
    return {"f1": f1, "precision": pr, "recall": rc}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", required=True)
    ap.add_argument("--data", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--epochs", type=float, default=4)
    ap.add_argument("--lr", type=float, default=3e-5)
    ap.add_argument("--batch", type=int, default=32)
    ap.add_argument("--max-len", type=int, default=384)
    ap.add_argument("--seed", type=int, default=42)
    # DeBERTa-v3 bf16'da ıraksıyor (14 Eki: loss 0.65 → 134 ilk 100 adımda); onun için --fp32
    ap.add_argument("--fp32", action="store_true")
    args = ap.parse_args()

    set_seed(args.seed)
    tok = AutoTokenizer.from_pretrained(args.base)
    data = Path(args.data)
    train = PinjData(data / "train.jsonl", tok, args.max_len)
    dev = PinjData(data / "dev.jsonl", tok, args.max_len)
    model = AutoModelForSequenceClassification.from_pretrained(
        args.base, num_labels=2, id2label=ID2LABEL, label2id=LABEL2ID)

    targs = TrainingArguments(
        output_dir=args.out, num_train_epochs=args.epochs, learning_rate=args.lr,
        per_device_train_batch_size=args.batch, per_device_eval_batch_size=64,
        warmup_steps=0.06, weight_decay=0.01, lr_scheduler_type="linear", bf16=not args.fp32,
        eval_strategy="epoch", save_strategy="epoch", save_total_limit=2, load_best_model_at_end=True,
        metric_for_best_model="f1", greater_is_better=True, logging_steps=50, report_to="none",
        seed=args.seed, train_sampling_strategy="group_by_length",
    )
    trainer = Trainer(model=model, args=targs, train_dataset=train, eval_dataset=dev,
                      data_collator=lambda b: collate(b, tok.pad_token_id), compute_metrics=compute_metrics)
    trainer.train()
    final = Path(args.out) / "final"
    trainer.save_model(final)
    # Checkpoint klasörlerinde tokenizer yok; final'e ayrıca yazılmazsa boş sözlük yüklenir (BSEBY dersi)
    tok.save_pretrained(final)
    hist = [h for h in trainer.state.log_history if "eval_f1" in h]
    rep = {"base": args.base, "data": args.data, "args": vars(args), "dev_by_epoch": hist,
           "best": trainer.state.best_metric}
    (Path(args.out) / "train_report.json").write_text(json.dumps(rep, indent=1), encoding="utf-8")
    for h in hist:
        print(f"epoch {h['epoch']:.0f}  dev F1 {h['eval_f1']:.4f}  P {h['eval_precision']:.4f}  "
              f"R {h['eval_recall']:.4f}  loss {h['eval_loss']:.4f}")


if __name__ == "__main__":
    main()