"""
Resume Section Classifier – Training Script

Fine-tunes distilbert-base-uncased for classifying resume text sections
into 8 categories: education, experience, skills, projects, summary,
certifications, contact, awards.

Author: Lorenzo Scaturchio (gr8monk3ys)

Usage:
    python train.py                           # Train with defaults
    python train.py --epochs 5 --batch-size 32
    python train.py --push-to-hub             # Push to HuggingFace Hub
    python train.py --output-dir ./my_model
"""

import json
import math
import logging
import sys
from pathlib import Path

import evaluate
import numpy as np
import torch
from datasets import DatasetDict
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    EarlyStoppingCallback,
    Trainer,
    TrainingArguments,
    set_seed,
)

from data_generator import generate_dataset, get_label_mapping, load_as_hf_dataset

# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
MODEL_NAME = "distilbert-base-uncased"
DEFAULT_OUTPUT_DIR = "./model_output"
DEFAULT_LOGGING_DIR = "./logs"
HUB_MODEL_ID = "gr8monk3ys/resume-section-classifier"
MAX_LENGTH = 256


# ---------------------------------------------------------------------------
# Metrics computation
# ---------------------------------------------------------------------------
def warmup_steps_for(
    n_train: int, batch_size: int, grad_accum: int, epochs: float, ratio: float
) -> int:
    """Convert a warm-up *ratio* into a step count.

    transformers 5.x removed ``warmup_ratio`` from TrainingArguments and kept
    only ``warmup_steps``; the v5 migration in 60a308b bumped the version
    without updating the call, so this script raised TypeError before reaching
    a single training step. The ratio stays the user-facing knob because it is
    the one that survives a change of batch size.
    """
    steps_per_epoch = math.ceil(n_train / max(batch_size * grad_accum, 1))
    return round(ratio * steps_per_epoch * epochs)


def build_compute_metrics():
    """Return a ``compute_metrics`` callable for the HF Trainer.

    Loads the four evaluate metrics once at creation time rather than per
    evaluation, to avoid repeated disk access.
    """
    accuracy_metric = evaluate.load("accuracy")
    f1_metric = evaluate.load("f1")
    precision_metric = evaluate.load("precision")
    recall_metric = evaluate.load("recall")

    def compute_metrics(eval_pred):
        logits, labels = eval_pred
        # Some model/config combinations return predictions as a tuple.
        if isinstance(logits, tuple):
            logits = logits[0]
        predictions = np.argmax(logits, axis=-1)

        acc = accuracy_metric.compute(predictions=predictions, references=labels)
        f1_macro = f1_metric.compute(
            predictions=predictions, references=labels, average="macro"
        )
        f1_weighted = f1_metric.compute(
            predictions=predictions, references=labels, average="weighted"
        )
        precision = precision_metric.compute(
            predictions=predictions, references=labels, average="weighted"
        )
        recall = recall_metric.compute(
            predictions=predictions, references=labels, average="weighted"
        )

        return {
            "accuracy": acc["accuracy"],
            "f1_macro": f1_macro["f1"],
            "f1_weighted": f1_weighted["f1"],
            "precision": precision["precision"],
            "recall": recall["recall"],
        }

    return compute_metrics


# ---------------------------------------------------------------------------
# Tokenization
# ---------------------------------------------------------------------------
def tokenize_dataset(
    dataset_dict: DatasetDict, tokenizer, label2id: dict, max_length: int = MAX_LENGTH
):
    """Tokenize all splits and encode labels as integers."""

    def preprocess(examples):
        tokenized = tokenizer(
            examples["text"],
            truncation=True,
            max_length=max_length,
            padding=False,  # Dynamic padding via DataCollator
        )
        tokenized["labels"] = [label2id[label] for label in examples["label"]]
        return tokenized

    tokenized = dataset_dict.map(
        preprocess,
        batched=True,
        remove_columns=["text", "label"],
        desc="Tokenizing",
    )

    return tokenized


# ---------------------------------------------------------------------------
# Main training function
# ---------------------------------------------------------------------------
def train(
    output_dir: str = DEFAULT_OUTPUT_DIR,
    model_name: str = MODEL_NAME,
    epochs: int = 4,
    batch_size: int = 16,
    learning_rate: float = 2e-5,
    weight_decay: float = 0.01,
    warmup_ratio: float = 0.1,
    max_length: int = MAX_LENGTH,
    examples_per_category: int = 80,
    augmented_copies: int = 2,
    seed: int = 42,
    push_to_hub: bool = False,
    hub_model_id: str = HUB_MODEL_ID,
    fp16: bool = None,
    gradient_accumulation_steps: int = 1,
    early_stopping_patience: int = 3,
):
    """
    Full training pipeline.

    Args:
        output_dir: Directory to save model and artifacts.
        model_name: Pretrained model identifier.
        epochs: Number of training epochs.
        batch_size: Training batch size.
        learning_rate: Peak learning rate.
        weight_decay: Weight decay for AdamW.
        warmup_ratio: Fraction of steps for warmup.
        max_length: Maximum token sequence length.
        examples_per_category: Base synthetic examples per category.
        augmented_copies: Augmented copies per base example.
        seed: Random seed.
        push_to_hub: Whether to push to HuggingFace Hub.
        hub_model_id: Hub model repository ID.
        fp16: Use mixed precision (auto-detected if None).
        gradient_accumulation_steps: Gradient accumulation steps.
        early_stopping_patience: Early stopping patience (epochs).
    """
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    # Auto-detect fp16
    if fp16 is None:
        fp16 = torch.cuda.is_available()

    logger.info("=" * 60)
    logger.info("Resume Section Classifier – Training")
    logger.info("=" * 60)
    logger.info(f"Model: {model_name}")
    logger.info(f"Output: {output_dir}")
    logger.info(f"Epochs: {epochs}, Batch size: {batch_size}, LR: {learning_rate}")
    logger.info(
        f"Device: {'CUDA' if torch.cuda.is_available() else 'MPS' if torch.backends.mps.is_available() else 'CPU'}"
    )
    logger.info(f"FP16: {fp16}")

    # Seed python, numpy and torch before anything samples from them.
    # TrainingArguments(seed=...) covers the Trainer, but not the synthetic data
    # generator below, which draws from the global `random` -- so without this
    # two runs at the same seed trained on different data.
    set_seed(seed)

    # ------------------------------------------------------------------
    # 1. Generate synthetic data
    # ------------------------------------------------------------------
    logger.info("\n[1/5] Generating synthetic training data...")
    raw_dataset = generate_dataset(
        examples_per_category=examples_per_category,
        augmented_copies=augmented_copies,
        seed=seed,
    )
    label2id, id2label = get_label_mapping(raw_dataset)
    num_labels = len(label2id)

    logger.info(f"  Total examples: {len(raw_dataset)}")
    logger.info(f"  Labels ({num_labels}): {list(label2id.keys())}")

    # Create HF DatasetDict with train/val/test splits
    dataset_dict = load_as_hf_dataset(raw_dataset)
    logger.info(f"  Train: {len(dataset_dict['train'])}")
    logger.info(f"  Validation: {len(dataset_dict['validation'])}")
    logger.info(f"  Test: {len(dataset_dict['test'])}")

    # ------------------------------------------------------------------
    # 2. Tokenize
    # ------------------------------------------------------------------
    logger.info("\n[2/5] Loading tokenizer and tokenizing data...")
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokenized_dataset = tokenize_dataset(dataset_dict, tokenizer, label2id, max_length)
    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

    # ------------------------------------------------------------------
    # 3. Load model
    # ------------------------------------------------------------------
    logger.info("\n[3/5] Loading pretrained model...")
    model = AutoModelForSequenceClassification.from_pretrained(
        model_name,
        num_labels=num_labels,
        id2label=id2label,
        label2id=label2id,
    )
    logger.info(f"  Parameters: {sum(p.numel() for p in model.parameters()):,}")
    logger.info(
        f"  Trainable: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}"
    )

    # ------------------------------------------------------------------
    # 4. Training
    # ------------------------------------------------------------------
    logger.info("\n[4/5] Training...")

    training_args = TrainingArguments(
        output_dir=output_dir,
        # Training hyperparameters
        num_train_epochs=epochs,
        per_device_train_batch_size=batch_size,
        per_device_eval_batch_size=batch_size * 2,
        gradient_accumulation_steps=gradient_accumulation_steps,
        learning_rate=learning_rate,
        weight_decay=weight_decay,
        warmup_steps=warmup_steps_for(
            len(tokenized_dataset["train"]),
            batch_size,
            gradient_accumulation_steps,
            epochs,
            warmup_ratio,
        ),
        lr_scheduler_type="linear",  # the library default; matches paper-classifier
        # Evaluation
        eval_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
        metric_for_best_model="f1_macro",
        greater_is_better=True,
        # Logging
        logging_strategy="steps",
        logging_steps=50,
        report_to="none",
        # Efficiency
        fp16=fp16,
        dataloader_num_workers=0,
        # Reproducibility
        seed=seed,
        data_seed=seed,
        # Hub
        push_to_hub=False,  # We'll push manually after evaluation
        # Misc
        save_total_limit=2,
        disable_tqdm=False,
    )

    callbacks = []
    if early_stopping_patience > 0:
        callbacks.append(
            EarlyStoppingCallback(early_stopping_patience=early_stopping_patience)
        )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset["train"],
        eval_dataset=tokenized_dataset["validation"],
        processing_class=tokenizer,
        data_collator=data_collator,
        compute_metrics=build_compute_metrics(),
        callbacks=callbacks,
    )

    train_result = trainer.train()

    # Log training metrics
    logger.info("\nTraining Results:")
    for key, value in train_result.metrics.items():
        logger.info(f"  {key}: {value}")

    # ------------------------------------------------------------------
    # 5. Evaluation
    # ------------------------------------------------------------------
    logger.info("\n[5/5] Evaluating on test set...")
    test_results = trainer.evaluate(tokenized_dataset["test"])

    logger.info("\nTest Results:")
    for key, value in test_results.items():
        logger.info(
            f"  {key}: {value:.4f}" if isinstance(value, float) else f"  {key}: {value}"
        )

    # ------------------------------------------------------------------
    # Save artifacts
    # ------------------------------------------------------------------
    logger.info("\nSaving model and artifacts...")

    # Save model + tokenizer
    final_path = output_path / "final_model"
    trainer.save_model(str(final_path))
    tokenizer.save_pretrained(str(final_path))

    # Save label mapping
    label_mapping = {
        "label2id": label2id,
        "id2label": {str(k): v for k, v in id2label.items()},
        "labels": list(label2id.keys()),
    }
    with open(final_path / "label_mapping.json", "w") as f:
        json.dump(label_mapping, f, indent=2)

    # Save training config
    train_config = {
        "model_name": model_name,
        "max_length": max_length,
        "epochs": epochs,
        "batch_size": batch_size,
        "learning_rate": learning_rate,
        "weight_decay": weight_decay,
        "warmup_ratio": warmup_ratio,
        "examples_per_category": examples_per_category,
        "augmented_copies": augmented_copies,
        "seed": seed,
        "num_labels": num_labels,
        "train_size": len(dataset_dict["train"]),
        "val_size": len(dataset_dict["validation"]),
        "test_size": len(dataset_dict["test"]),
    }
    with open(final_path / "training_config.json", "w") as f:
        json.dump(train_config, f, indent=2)

    # Save metrics
    all_metrics = {
        "train": train_result.metrics,
        "test": test_results,
    }
    with open(final_path / "metrics.json", "w") as f:
        json.dump(all_metrics, f, indent=2)

    logger.info(f"\nAll artifacts saved to: {final_path}")

    # ------------------------------------------------------------------
    # Optional: Push to Hub
    # ------------------------------------------------------------------
    if push_to_hub:
        logger.info(f"\nPushing to HuggingFace Hub: {hub_model_id}")
        try:
            trainer.push_to_hub(
                repo_id=hub_model_id,
                commit_message="Upload fine-tuned resume section classifier",
            )
            tokenizer.push_to_hub(hub_model_id)
            logger.info("Successfully pushed to Hub!")
        except Exception as exc:  # noqa: BLE001 - reported, then fatal
            logger.error("Failed to push to Hub: %s", exc)
            logger.info("You can push manually later with:")
            logger.info("  huggingface-cli upload %s %s", hub_model_id, final_path)
            # Fatal, as in paper-classifier: a run whose artifact never reached
            # the Hub has not succeeded, and exiting 0 hides that from CI.
            sys.exit(1)

    logger.info("\nTraining complete!")
    return test_results


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description="Fine-tune DistilBERT for resume section classification",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    # Model & output
    parser.add_argument(
        "--model-name",
        type=str,
        default=MODEL_NAME,
        help="Pretrained model name or path",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default=DEFAULT_OUTPUT_DIR,
        help="Output directory for model and artifacts",
    )

    # Training hyperparameters
    parser.add_argument(
        "--epochs", type=int, default=4, help="Number of training epochs"
    )
    parser.add_argument(
        "--batch-size", type=int, default=16, help="Training batch size per device"
    )
    parser.add_argument(
        "--learning-rate", type=float, default=2e-5, help="Peak learning rate"
    )
    parser.add_argument(
        "--weight-decay", type=float, default=0.01, help="Weight decay for AdamW"
    )
    parser.add_argument(
        "--warmup-ratio",
        type=float,
        default=0.1,
        help="Fraction of total steps for linear warmup",
    )
    parser.add_argument(
        "--max-length",
        type=int,
        default=MAX_LENGTH,
        help="Maximum token sequence length",
    )
    parser.add_argument(
        "--gradient-accumulation-steps",
        type=int,
        default=1,
        help="Number of gradient accumulation steps",
    )

    # Data
    parser.add_argument(
        "--examples-per-category",
        type=int,
        default=80,
        help="Base synthetic examples per category",
    )
    parser.add_argument(
        "--augmented-copies",
        type=int,
        default=2,
        help="Augmented copies per base example",
    )
    parser.add_argument(
        "--seed", type=int, default=42, help="Random seed for reproducibility"
    )

    # Training config
    parser.add_argument(
        "--fp16", action="store_true", default=None, help="Force FP16 training"
    )
    parser.add_argument("--no-fp16", action="store_true", help="Disable FP16 training")
    parser.add_argument(
        "--early-stopping-patience",
        type=int,
        default=3,
        help="Early stopping patience (0 to disable)",
    )

    # Hub
    parser.add_argument(
        "--push-to-hub",
        action="store_true",
        help="Push trained model to HuggingFace Hub",
    )
    parser.add_argument(
        "--hub-model-id",
        type=str,
        default=HUB_MODEL_ID,
        help="HuggingFace Hub model ID",
    )

    args = parser.parse_args()

    # Handle fp16 flags
    fp16 = args.fp16
    if args.no_fp16:
        fp16 = False

    results = train(
        output_dir=args.output_dir,
        model_name=args.model_name,
        epochs=args.epochs,
        batch_size=args.batch_size,
        learning_rate=args.learning_rate,
        weight_decay=args.weight_decay,
        warmup_ratio=args.warmup_ratio,
        max_length=args.max_length,
        examples_per_category=args.examples_per_category,
        augmented_copies=args.augmented_copies,
        seed=args.seed,
        push_to_hub=args.push_to_hub,
        hub_model_id=args.hub_model_id,
        fp16=fp16,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        early_stopping_patience=args.early_stopping_patience,
    )