"""
PyTorch Lightning integration for Armory SDK.

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

from typing import Any, Dict, Optional

import armory

# Lazy import
try:
    from pytorch_lightning.loggers import Logger
    from pytorch_lightning.utilities.rank_zero import rank_zero_only

    HAS_LIGHTNING = True
except ImportError:
    HAS_LIGHTNING = False
    Logger = object

    def rank_zero_only(fn):
        return fn


class ArmoryLightningLogger(Logger):
    """PyTorch Lightning logger that sends metrics to Armory.

    Usage:
        import armory
        from armory.integrations import ArmoryLightningLogger

        with armory.start_session(project="my-project", goal="train model"):
            logger = ArmoryLightningLogger()
            trainer = Trainer(logger=logger)
            trainer.fit(model, datamodule)
    """

    def __init__(
        self,
        name: str = "armory",
        version: Optional[str] = None,
    ):
        if not HAS_LIGHTNING:
            raise ImportError(
                "pytorch_lightning is required for ArmoryLightningLogger. "
                "Install it with: pip install pytorch-lightning"
            )

        super().__init__()
        self._name = name
        self._version = version

    @property
    def name(self) -> str:
        return self._name

    @property
    def version(self) -> Optional[str]:
        return self._version

    @rank_zero_only
    def log_hyperparams(self, params: Dict[str, Any]) -> None:
        """Log hyperparameters."""
        # Convert to flat dict with string keys
        flat_params = {}
        for key, value in params.items():
            if isinstance(value, (int, float, str, bool)):
                flat_params[key] = value
            else:
                flat_params[key] = str(value)

        armory.log_params(flat_params)

    @rank_zero_only
    def log_metrics(
        self,
        metrics: Dict[str, float],
        step: Optional[int] = None,
    ) -> None:
        """Log metrics."""
        # Filter to numeric values
        numeric_metrics = {
            k: v for k, v in metrics.items()
            if isinstance(v, (int, float))
        }

        if numeric_metrics:
            armory.log_metrics(numeric_metrics, step=step)

    @rank_zero_only
    def save(self) -> None:
        """Save is not needed for Armory (server-side storage)."""
        pass

    @rank_zero_only
    def finalize(self, status: str) -> None:
        """Finalize logging."""
        # Status is automatically handled by session end
        pass
