"""
Armory Callbacks - Training framework integrations

This module provides callbacks for popular training frameworks:
- PyTorch Lightning
- HuggingFace Transformers Trainer
- Generic training loops
"""

from typing import Any, Dict, Optional
from armory.run import Run, init
from armory.logger import get_logger

logger = get_logger("callback")


class ArmoryCallback:
    """
    Base callback for Armory logging.
    
    Can be used with any training loop:
    
        callback = ArmoryCallback(project="my-project")
        callback.on_train_begin()
        
        for epoch in range(num_epochs):
            callback.on_epoch_begin(epoch)
            for batch in dataloader:
                loss = train_step(batch)
                callback.on_batch_end(loss=loss)
            callback.on_epoch_end(epoch, metrics={"val_loss": val_loss})
        
        callback.on_train_end()
    """
    
    def __init__(
        self,
        project: str = "default",
        name: Optional[str] = None,
        config: Optional[Dict[str, Any]] = None,
        run: Optional[Run] = None,
        **kwargs
    ):
        self._run = run
        self._project = project
        self._name = name
        self._config = config or {}
        self._kwargs = kwargs
        self._step = 0
    
    @property
    def run(self) -> Run:
        if self._run is None:
            self._run = init(
                project=self._project,
                name=self._name,
                config=self._config,
                **self._kwargs
            )
        return self._run
    
    def on_train_begin(self, **kwargs):
        """Called at the start of training"""
        self.run.log_event("train_start", "Training started")
    
    def on_train_end(self, **kwargs):
        """Called at the end of training"""
        self.run.log_event("train_end", "Training completed")
        self.run.finish()
    
    def on_epoch_begin(self, epoch: int, **kwargs):
        """Called at the start of each epoch"""
        self.run.log_event("epoch_start", f"Epoch {epoch} started")
    
    def on_epoch_end(self, epoch: int, metrics: Optional[Dict[str, float]] = None, **kwargs):
        """Called at the end of each epoch"""
        if metrics:
            self.run.log(metrics)
        self.run.log_event("epoch_end", f"Epoch {epoch} completed", metrics=metrics)
    
    def on_batch_end(self, step: Optional[int] = None, **metrics):
        """Called at the end of each batch"""
        if step is not None:
            self._step = step
        else:
            self._step += 1
        
        if metrics:
            self.run.log(metrics, step=self._step)
    
    def on_evaluate_begin(self, **kwargs):
        """Called at the start of evaluation"""
        self.run.log_event("eval_start", "Evaluation started")
    
    def on_evaluate_end(self, metrics: Optional[Dict[str, float]] = None, **kwargs):
        """Called at the end of evaluation"""
        if metrics:
            self.run.log(metrics)
        self.run.log_event("eval_end", "Evaluation completed", metrics=metrics)
    
    def on_save(self, path: str = "", **kwargs):
        """Called when a checkpoint is saved"""
        self.run.log_event("checkpoint", f"Checkpoint saved: {path}")
    
    def on_error(self, error: str, **kwargs):
        """Called when an error occurs"""
        self.run.log_event("error", error)


# PyTorch Lightning integration
try:
    from pytorch_lightning.callbacks import Callback as PLCallback
    
    class ArmoryLightningCallback(PLCallback, ArmoryCallback):
        """
        PyTorch Lightning callback for Armory.
        
        Usage:
            from armory import ArmoryLightningCallback
            
            trainer = pl.Trainer(
                callbacks=[ArmoryLightningCallback(project="my-project")]
            )
        """
        
        def __init__(self, **kwargs):
            PLCallback.__init__(self)
            ArmoryCallback.__init__(self, **kwargs)
        
        def on_train_start(self, trainer, pl_module):
            self.on_train_begin()
        
        def on_train_end(self, trainer, pl_module):
            ArmoryCallback.on_train_end(self)
        
        def on_train_epoch_start(self, trainer, pl_module):
            self.on_epoch_begin(trainer.current_epoch)
        
        def on_train_epoch_end(self, trainer, pl_module):
            metrics = {k: v.item() if hasattr(v, 'item') else v 
                      for k, v in trainer.callback_metrics.items()}
            self.on_epoch_end(trainer.current_epoch, metrics=metrics)
        
        def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
            if isinstance(outputs, dict) and 'loss' in outputs:
                loss = outputs['loss']
                if hasattr(loss, 'item'):
                    loss = loss.item()
                self.on_batch_end(step=trainer.global_step, loss=loss)
        
        def on_validation_start(self, trainer, pl_module):
            self.on_evaluate_begin()
        
        def on_validation_end(self, trainer, pl_module):
            metrics = {k: v.item() if hasattr(v, 'item') else v 
                      for k, v in trainer.callback_metrics.items()
                      if 'val' in k}
            self.on_evaluate_end(metrics=metrics)
        
        def on_save_checkpoint(self, trainer, pl_module, checkpoint):
            self.on_save(path=trainer.checkpoint_callback.best_model_path if trainer.checkpoint_callback else "")

except ImportError:
    # PyTorch Lightning not installed
    pass


# HuggingFace Transformers integration
try:
    from transformers import TrainerCallback
    
    class ArmoryTrainerCallback(TrainerCallback, ArmoryCallback):
        """
        HuggingFace Transformers Trainer callback for Armory.
        
        Usage:
            from armory import ArmoryTrainerCallback
            
            trainer = Trainer(
                model=model,
                args=training_args,
                callbacks=[ArmoryTrainerCallback(project="my-project")]
            )
        """
        
        def __init__(self, **kwargs):
            TrainerCallback.__init__(self)
            ArmoryCallback.__init__(self, **kwargs)
        
        def on_train_begin(self, args, state, control, **kwargs):
            logger.info(f"Training started, run_id: {self.run.id}")
            ArmoryCallback.on_train_begin(self)

        def on_train_end(self, args, state, control, **kwargs):
            logger.info(f"Training ended, finalizing run {self.run.id}...")
            ArmoryCallback.on_train_end(self)
            logger.info("Training finalization complete")
        
        def on_epoch_begin(self, args, state, control, **kwargs):
            ArmoryCallback.on_epoch_begin(self, int(state.epoch))
        
        def on_epoch_end(self, args, state, control, **kwargs):
            ArmoryCallback.on_epoch_end(self, int(state.epoch))
        
        def on_log(self, args, state, control, logs=None, **kwargs):
            if logs:
                metrics = {k: float(v) for k, v in logs.items() if isinstance(v, (int, float))}
                if metrics:
                    logger.debug(f"Logging metrics at step {state.global_step}: {list(metrics.keys())}")
                    self.run.log(metrics, step=state.global_step)
                else:
                    logger.debug(f"on_log called but no numeric metrics in logs: {list(logs.keys())}")
        
        def on_evaluate(self, args, state, control, metrics=None, **kwargs):
            self.on_evaluate_end(metrics=metrics)
        
        def on_save(self, args, state, control, **kwargs):
            ArmoryCallback.on_save(self, path=args.output_dir)

except ImportError:
    # Transformers not installed
    pass
