"""
HuggingFace Transformers integration for Armory SDK.

Provides automatic logging of training metrics, hyperparameters,
and checkpoints to Armory.
"""

from typing import Any, Dict, Optional

import armory

# Lazy import to avoid dependency if not installed
try:
    from transformers import TrainerCallback, TrainerControl, TrainerState, TrainingArguments

    HAS_TRANSFORMERS = True
except ImportError:
    HAS_TRANSFORMERS = False
    TrainerCallback = object


class ArmoryTrainerCallback(TrainerCallback):
    """HuggingFace Trainer callback that automatically logs to Armory.

    This callback:
    - Logs all training metrics (loss, accuracy, etc.)
    - Logs evaluation metrics
    - Records checkpoint saves as artifacts
    - Handles errors gracefully

    Usage:
        import armory
        from armory.integrations import ArmoryTrainerCallback

        with armory.start_session(project="my-project", goal="finetune bert"):
            trainer = Trainer(
                model=model,
                args=training_args,
                callbacks=[ArmoryTrainerCallback()],
            )
            trainer.train()
    """

    def __init__(
        self,
        log_checkpoints: bool = True,
        log_model_config: bool = True,
    ):
        """Initialize the callback.

        Args:
            log_checkpoints: Whether to log checkpoint saves as artifacts
            log_model_config: Whether to log model config as parameters
        """
        if not HAS_TRANSFORMERS:
            raise ImportError(
                "transformers is required for ArmoryTrainerCallback. "
                "Install it with: pip install transformers"
            )

        self.log_checkpoints = log_checkpoints
        self.log_model_config = log_model_config
        self._logged_config = False

    def on_train_begin(
        self,
        args: "TrainingArguments",
        state: "TrainerState",
        control: "TrainerControl",
        model=None,
        **kwargs,
    ):
        """Log training configuration at the start."""
        # Log training arguments as parameters
        armory.log_params({
            "learning_rate": args.learning_rate,
            "batch_size": args.per_device_train_batch_size,
            "gradient_accumulation_steps": args.gradient_accumulation_steps,
            "num_train_epochs": args.num_train_epochs,
            "max_steps": args.max_steps,
            "warmup_steps": args.warmup_steps,
            "warmup_ratio": args.warmup_ratio,
            "weight_decay": args.weight_decay,
            "adam_epsilon": args.adam_epsilon,
            "max_grad_norm": args.max_grad_norm,
            "seed": args.seed,
            "fp16": args.fp16,
            "bf16": args.bf16,
        })

        # Log model config if available
        if self.log_model_config and model is not None and not self._logged_config:
            try:
                config = model.config.to_dict()
                # Only log key config values to avoid too many params
                key_params = [
                    "hidden_size", "num_hidden_layers", "num_attention_heads",
                    "vocab_size", "max_position_embeddings", "model_type",
                ]
                model_params = {
                    f"model_{k}": config[k]
                    for k in key_params
                    if k in config
                }
                armory.log_params(model_params)
                self._logged_config = True
            except Exception:
                pass  # Ignore if model config is not available

    def on_log(
        self,
        args: "TrainingArguments",
        state: "TrainerState",
        control: "TrainerControl",
        logs: Optional[Dict[str, Any]] = None,
        **kwargs,
    ):
        """Log metrics at each logging step."""
        if logs is None:
            return

        # Filter to only numeric values
        metrics = {}
        for key, value in logs.items():
            if isinstance(value, (int, float)):
                # Skip epoch as we log it separately
                if key != "epoch":
                    metrics[key] = value

        if metrics:
            armory.log_metrics(metrics, step=state.global_step)

    def on_epoch_end(
        self,
        args: "TrainingArguments",
        state: "TrainerState",
        control: "TrainerControl",
        **kwargs,
    ):
        """Log at epoch end."""
        if state.epoch is not None:
            armory.log_metric("epoch", state.epoch, step=state.global_step)

    def on_evaluate(
        self,
        args: "TrainingArguments",
        state: "TrainerState",
        control: "TrainerControl",
        metrics: Optional[Dict[str, float]] = None,
        **kwargs,
    ):
        """Log evaluation metrics."""
        if metrics:
            # Filter to numeric values
            eval_metrics = {
                k: v for k, v in metrics.items()
                if isinstance(v, (int, float))
            }
            if eval_metrics:
                armory.log_metrics(eval_metrics, step=state.global_step)

    def on_save(
        self,
        args: "TrainingArguments",
        state: "TrainerState",
        control: "TrainerControl",
        **kwargs,
    ):
        """Log checkpoint saves."""
        if self.log_checkpoints:
            checkpoint_path = f"{args.output_dir}/checkpoint-{state.global_step}"
            armory.log_artifact(
                name=f"checkpoint-{state.global_step}",
                path=checkpoint_path,
                artifact_type="model",
            )

    def on_train_end(
        self,
        args: "TrainingArguments",
        state: "TrainerState",
        control: "TrainerControl",
        **kwargs,
    ):
        """Log final training stats."""
        armory.log_metrics({
            "total_steps": state.global_step,
            "total_epochs": state.epoch or 0,
        })
