"""
Fine-Tune Custom Email Classifiers — Google Colab Training Script
=================================================================

Copy this entire file into a Google Colab notebook (one cell per section)
or run it as a standalone Python script.

Trains two binary classifiers on the project's JSONL datasets:
  1. **Malicious Intent** — phishing / fraud detection
  2. **Prompt Injection** — LLM injection attack detection

Recommended base models (all < 100M params, sub-1s CPU inference):
  • microsoft/MiniLM-L12-H384-uncased          (33M params)
  • microsoft/deberta-v3-small                  (44M params)
  • distilbert-base-uncased                     (66M params)

After training, the script exports the model + tokenizer to a directory
that can be copied into the project at ``models/custom_<task>/`` for
use with ``CustomL2SemanticAnalyzer``.

Usage (Colab):
    1. Upload train.jsonl, val.jsonl, test.jsonl for each task.
    2. Run all cells.
    3. Download the exported model directories.
    4. Place them in the project under ``models/``.
"""

# ============================================================
# SECTION 1: Setup & Installs
# ============================================================

# !pip install -q transformers datasets accelerate safetensors scikit-learn

import json
import os
from pathlib import Path

import numpy as np
import torch
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    f1_score,
    precision_score,
    recall_score,
)

from datasets import Dataset
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    Trainer,
    TrainingArguments,
    EarlyStoppingCallback,
)


# ============================================================
# SECTION 2: Configuration
# ============================================================

class TrainingConfig:
    """Central configuration for the training run.

    Edit these values before running. All paths are relative to the
    Colab working directory (or script directory).
    """

    # --- Model ---
    # Choose one of:
    #   "microsoft/MiniLM-L12-H384-uncased"     (33M, fastest)
    #   "microsoft/deberta-v3-small"             (44M, best quality)
    #   "distilbert-base-uncased"                (66M, good balance)
    BASE_MODEL: str = "microsoft/deberta-v3-small"

    # --- Task ---
    # Set to "malicious_intent" or "prompt_injection"
    TASK: str = "malicious_intent"

    # --- Data paths ---
    TRAIN_PATH: str = f"data/{TASK}/train.jsonl"
    VAL_PATH: str = f"data/{TASK}/val.jsonl"
    TEST_PATH: str = f"data/{TASK}/test.jsonl"

    # --- Tokenizer ---
    MAX_LENGTH: int = 256

    # --- Training hyperparameters ---
    EPOCHS: int = 5
    BATCH_SIZE: int = 16
    LEARNING_RATE: float = 2e-5
    WEIGHT_DECAY: float = 0.01
    WARMUP_RATIO: float = 0.1
    FP16: bool = torch.cuda.is_available()

    # --- Early stopping ---
    EARLY_STOPPING_PATIENCE: int = 2

    # --- Output ---
    OUTPUT_DIR: str = f"output/{TASK}"
    EXPORT_DIR: str = f"export/custom_{TASK}"

    # --- Reproducibility ---
    SEED: int = 42


cfg = TrainingConfig()

# Set seeds for reproducibility
torch.manual_seed(cfg.SEED)
np.random.seed(cfg.SEED)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(cfg.SEED)


# ============================================================
# SECTION 3: Data Loading
# ============================================================

def load_jsonl_dataset(path: str) -> list[dict]:
    """Load a JSONL file matching the project's dataset schema.

    Expected fields per line:
        - text_body (str): Plain-text email body.
        - html_body (str): HTML email body (used as fallback).
        - label (int): Binary label (0 = benign, 1 = malicious/injection).
        - id (str, optional): Sample identifier.
    """
    samples = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            # Use text_body if available, fall back to html_body
            text = row.get("text_body", "") or ""
            if not text:
                text = row.get("html_body", "") or ""
            samples.append({
                "text": text,
                "label": int(row.get("label", 0)),
                "id": row.get("id", ""),
            })
    return samples


def prepare_datasets() -> tuple[Dataset, Dataset, Dataset | None]:
    """Load train/val/test splits and convert to HuggingFace Datasets."""
    train_data = load_jsonl_dataset(cfg.TRAIN_PATH)
    val_data = load_jsonl_dataset(cfg.VAL_PATH)

    test_data = None
    if os.path.exists(cfg.TEST_PATH):
        test_data = load_jsonl_dataset(cfg.TEST_PATH)

    print(f"Task:       {cfg.TASK}")
    print(f"Train:      {len(train_data)} samples")
    print(f"Validation: {len(val_data)} samples")
    if test_data:
        print(f"Test:       {len(test_data)} samples")

    # Label distribution
    train_pos = sum(1 for s in train_data if s["label"] == 1)
    print(f"\nTrain label distribution: "
          f"{train_pos} positive ({train_pos/len(train_data)*100:.1f}%), "
          f"{len(train_data)-train_pos} negative "
          f"({(len(train_data)-train_pos)/len(train_data)*100:.1f}%)")

    train_ds = Dataset.from_list(train_data)
    val_ds = Dataset.from_list(val_data)
    test_ds = Dataset.from_list(test_data) if test_data else None

    return train_ds, val_ds, test_ds


train_ds, val_ds, test_ds = prepare_datasets()


# ============================================================
# SECTION 4: Tokenization
# ============================================================

tokenizer = AutoTokenizer.from_pretrained(cfg.BASE_MODEL)


def tokenize_function(examples: dict) -> dict:
    """Tokenize the 'text' field with truncation and padding."""
    return tokenizer(
        examples["text"],
        truncation=True,
        max_length=cfg.MAX_LENGTH,
        padding="max_length",
    )


print(f"\nTokenizing with {cfg.BASE_MODEL} (max_length={cfg.MAX_LENGTH})...")

train_ds_tok = train_ds.map(tokenize_function, batched=True, batch_size=1000)
val_ds_tok = val_ds.map(tokenize_function, batched=True, batch_size=1000)
test_ds_tok = test_ds.map(tokenize_function, batched=True, batch_size=1000) if test_ds else None

# Set format for PyTorch
columns = ["input_ids", "attention_mask", "label"]
if "token_type_ids" in train_ds_tok.column_names:
    columns.append("token_type_ids")

train_ds_tok.set_format("torch", columns=columns)
val_ds_tok.set_format("torch", columns=columns)
if test_ds_tok:
    test_ds_tok.set_format("torch", columns=columns)

print("Tokenization complete.")


# ============================================================
# SECTION 5: Model Setup
# ============================================================

model = AutoModelForSequenceClassification.from_pretrained(
    cfg.BASE_MODEL,
    num_labels=2,
    id2label={0: "BENIGN", 1: "MALICIOUS"},
    label2id={"BENIGN": 0, "MALICIOUS": 1},
)

param_count = sum(p.numel() for p in model.parameters())
trainable_count = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"\nModel: {cfg.BASE_MODEL}")
print(f"Total parameters:     {param_count:>12,}")
print(f"Trainable parameters: {trainable_count:>12,}")
print(f"Model size:           {param_count * 4 / 1e6:>10.1f} MB (FP32 est.)")


# ============================================================
# SECTION 6: Metrics
# ============================================================

def compute_metrics(eval_pred) -> dict:
    """Compute precision, recall, F1, and accuracy for the Trainer."""
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return {
        "accuracy": accuracy_score(labels, predictions),
        "precision": precision_score(labels, predictions, zero_division=0),
        "recall": recall_score(labels, predictions, zero_division=0),
        "f1": f1_score(labels, predictions, zero_division=0),
    }


# ============================================================
# SECTION 7: Training
# ============================================================

training_args = TrainingArguments(
    output_dir=cfg.OUTPUT_DIR,
    num_train_epochs=cfg.EPOCHS,
    per_device_train_batch_size=cfg.BATCH_SIZE,
    per_device_eval_batch_size=cfg.BATCH_SIZE * 2,
    learning_rate=cfg.LEARNING_RATE,
    weight_decay=cfg.WEIGHT_DECAY,
    warmup_ratio=cfg.WARMUP_RATIO,
    fp16=cfg.FP16,
    # Evaluation
    eval_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    greater_is_better=True,
    # Logging
    logging_steps=50,
    logging_first_step=True,
    report_to="none",
    # Reproducibility
    seed=cfg.SEED,
    data_seed=cfg.SEED,
    # Save disk space
    save_total_limit=2,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_ds_tok,
    eval_dataset=val_ds_tok,
    compute_metrics=compute_metrics,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=cfg.EARLY_STOPPING_PATIENCE)],
)

print(f"\n{'='*60}")
print(f"Starting training: {cfg.TASK}")
print(f"Base model: {cfg.BASE_MODEL}")
print(f"Epochs: {cfg.EPOCHS}, Batch size: {cfg.BATCH_SIZE}")
print(f"Learning rate: {cfg.LEARNING_RATE}, FP16: {cfg.FP16}")
print(f"{'='*60}\n")

trainer.train()


# ============================================================
# SECTION 8: Evaluation
# ============================================================

print(f"\n{'='*60}")
print("Validation Set Evaluation")
print(f"{'='*60}")

val_results = trainer.evaluate(val_ds_tok)
for key, value in sorted(val_results.items()):
    if isinstance(value, float):
        print(f"  {key}: {value:.4f}")

if test_ds_tok:
    print(f"\n{'='*60}")
    print("Test Set Evaluation")
    print(f"{'='*60}")

    test_results = trainer.evaluate(test_ds_tok)
    for key, value in sorted(test_results.items()):
        if isinstance(value, float):
            print(f"  {key}: {value:.4f}")

    # Detailed classification report
    test_pred = trainer.predict(test_ds_tok)
    test_preds = np.argmax(test_pred.predictions, axis=-1)
    print("\nDetailed Classification Report:")
    print(classification_report(
        test_pred.label_ids,
        test_preds,
        target_names=["BENIGN", "MALICIOUS"],
    ))


# ============================================================
# SECTION 9: Threshold Tuning
# ============================================================

def find_optimal_threshold(
    logits: np.ndarray,
    labels: np.ndarray,
    recall_floor: float = 0.85,
) -> tuple[float, dict]:
    """Sweep thresholds on the positive-class probability.

    Finds the threshold that maximizes F1 while maintaining recall
    above the specified floor. This mirrors the grid search logic in
    ``scripts/evaluate_pipeline.py``.

    Args:
        logits: Raw model logits (N, 2).
        labels: Ground-truth binary labels (N,).
        recall_floor: Minimum recall requirement.

    Returns:
        Tuple of (best_threshold, metrics_at_threshold).
    """
    probs = torch.softmax(torch.tensor(logits), dim=-1)[:, 1].numpy()

    best_threshold = 0.5
    best_f1 = -1.0
    best_metrics = {}

    for t in np.arange(0.05, 0.96, 0.01):
        preds = (probs >= t).astype(int)
        prec = precision_score(labels, preds, zero_division=0)
        rec = recall_score(labels, preds, zero_division=0)
        f1 = f1_score(labels, preds, zero_division=0)

        if rec >= recall_floor and f1 > best_f1:
            best_f1 = f1
            best_threshold = float(t)
            best_metrics = {
                "threshold": float(t),
                "precision": float(prec),
                "recall": float(rec),
                "f1": float(f1),
            }

    # Fallback: if no threshold meets recall floor, take max F1
    if best_f1 < 0:
        for t in np.arange(0.05, 0.96, 0.01):
            preds = (probs >= t).astype(int)
            f1 = f1_score(labels, preds, zero_division=0)
            if f1 > best_f1:
                best_f1 = f1
                best_threshold = float(t)
                prec = precision_score(labels, preds, zero_division=0)
                rec = recall_score(labels, preds, zero_division=0)
                best_metrics = {
                    "threshold": float(t),
                    "precision": float(prec),
                    "recall": float(rec),
                    "f1": float(f1),
                }

    return best_threshold, best_metrics


print(f"\n{'='*60}")
print("Threshold Tuning (on validation set)")
print(f"{'='*60}")

val_pred = trainer.predict(val_ds_tok)
optimal_t, optimal_m = find_optimal_threshold(
    val_pred.predictions, val_pred.label_ids, recall_floor=0.85,
)
print(f"\n  Optimal threshold:  {optimal_t:.2f}")
print(f"  Precision:          {optimal_m.get('precision', 0):.4f}")
print(f"  Recall:             {optimal_m.get('recall', 0):.4f}")
print(f"  F1:                 {optimal_m.get('f1', 0):.4f}")


# ============================================================
# SECTION 10: Export Model
# ============================================================

export_path = Path(cfg.EXPORT_DIR)
export_path.mkdir(parents=True, exist_ok=True)

# Save model + tokenizer
trainer.save_model(str(export_path))
tokenizer.save_pretrained(str(export_path))

# Save training metadata
metadata = {
    "task": cfg.TASK,
    "base_model": cfg.BASE_MODEL,
    "max_length": cfg.MAX_LENGTH,
    "optimal_threshold": optimal_t,
    "optimal_metrics": optimal_m,
    "val_metrics": {
        k: v for k, v in val_results.items()
        if isinstance(v, (int, float))
    },
    "training_config": {
        "epochs": cfg.EPOCHS,
        "batch_size": cfg.BATCH_SIZE,
        "learning_rate": cfg.LEARNING_RATE,
        "weight_decay": cfg.WEIGHT_DECAY,
        "warmup_ratio": cfg.WARMUP_RATIO,
        "seed": cfg.SEED,
    },
    "param_count": param_count,
}

if test_ds_tok:
    metadata["test_metrics"] = {
        k: v for k, v in test_results.items()
        if isinstance(v, (int, float))
    }

with open(export_path / "training_metadata.json", "w") as f:
    json.dump(metadata, f, indent=2)

print(f"\n{'='*60}")
print(f"Model exported to: {export_path}")
print(f"{'='*60}")
print(f"\nFiles in export directory:")
for p in sorted(export_path.iterdir()):
    size_kb = p.stat().st_size / 1024
    print(f"  {p.name:40s} {size_kb:>8.1f} KB")


# ============================================================
# SECTION 11: Integration Guide
# ============================================================

print(f"""
{'='*60}
INTEGRATION GUIDE
{'='*60}

1. Copy the exported directory to your project:

   cp -r {export_path} /path/to/Malicious-Email-Scorer/models/custom_{cfg.TASK}/

2. In your app/main.py lifespan(), replace the L2 loader:

   # BEFORE (off-the-shelf):
   # from app.engines.semantic.orchestrator import load_models, shutdown_models
   # load_models()

   # AFTER (custom):
   from app.engines.semantic.custom_adapter import (
       CustomL2SemanticAnalyzer,
       load_custom_models,
       shutdown_custom_models,
   )
   load_custom_models(
       malicious_threshold={optimal_t:.2f},  # from threshold tuning
   )

3. Register the custom analyzer instead of the OTS one:

   manager = AnalysisManager()
   manager.register(L1HeuristicsAnalyzer())
   manager.register(CustomL2SemanticAnalyzer())  # <-- custom

4. Run the evaluation script to tune pipeline thresholds:

   python scripts/evaluate_pipeline.py --run-test-eval

5. Recommended threshold for {cfg.TASK}: {optimal_t:.2f}
   (F1={optimal_m.get('f1', 0):.4f}, Recall={optimal_m.get('recall', 0):.4f})
""")