"""Global logging functions for Armory SDK.

These functions operate on the currently active session,
similar to MLflow's global logging API.

Usage:
    import armory

    with armory.start_session(project="my-project", goal="train"):
        armory.log_param("learning_rate", 0.001)
        armory.log_metric("loss", 0.5, step=1)
        armory.set_tag("experiment_type", "baseline")
"""

from typing import Any, Dict, Optional, Union

from armory.session import get_active_session


def _get_session_or_raise():
    """Get active session or raise error."""
    session = get_active_session()
    if session is None:
        raise RuntimeError(
            "No active Armory session. Use `with armory.start_session(...)` first."
        )
    return session


def log_param(key: str, value: Any) -> None:
    """Log a parameter to the active session.

    Args:
        key: Parameter name
        value: Parameter value (will be converted to string for storage)

    Example:
        armory.log_param("learning_rate", 0.001)
        armory.log_param("batch_size", 32)
        armory.log_param("optimizer", "adamw")
    """
    session = _get_session_or_raise()
    session.log_param(key, value)


def log_params(params: Dict[str, Any]) -> None:
    """Log multiple parameters to the active session.

    Args:
        params: Dictionary of parameter names to values

    Example:
        armory.log_params({
            "learning_rate": 0.001,
            "batch_size": 32,
            "optimizer": "adamw",
            "warmup_steps": 100,
        })
    """
    session = _get_session_or_raise()
    session.log_params(params)


def log_metric(
    key: str,
    value: float,
    step: Optional[int] = None,
) -> None:
    """Log a metric to the active session.

    Args:
        key: Metric name
        value: Metric value
        step: Optional step number (e.g., epoch, iteration)

    Example:
        # Log without step
        armory.log_metric("final_accuracy", 0.95)

        # Log with step (for tracking over time)
        for epoch in range(10):
            loss = train_epoch()
            armory.log_metric("loss", loss, step=epoch)
    """
    session = _get_session_or_raise()
    session.log_metric(key, value, step)


def log_metrics(
    metrics: Dict[str, float],
    step: Optional[int] = None,
) -> None:
    """Log multiple metrics to the active session.

    Args:
        metrics: Dictionary of metric names to values
        step: Optional step number

    Example:
        armory.log_metrics({
            "train_loss": 0.5,
            "train_accuracy": 0.92,
            "val_loss": 0.6,
            "val_accuracy": 0.89,
        }, step=epoch)
    """
    session = _get_session_or_raise()
    session.log_metrics(metrics, step)


def log_artifact(
    name: str,
    path: str,
    artifact_type: str = "file",
) -> None:
    """Log an artifact to the active session.

    Args:
        name: Artifact name
        path: Path to the artifact file
        artifact_type: Type of artifact ("file", "model", "dataset", "config")

    Example:
        armory.log_artifact("model_checkpoint", "./checkpoints/best.pt", "model")
        armory.log_artifact("training_config", "./config.yaml", "config")
    """
    session = _get_session_or_raise()
    session.log_artifact(name, path, artifact_type)


def log_error(error: Union[str, Exception]) -> None:
    """Log an error to the active session.

    The error will be included in the retrospective analysis.

    Args:
        error: Error message or exception

    Example:
        try:
            train()
        except Exception as e:
            armory.log_error(e)
            raise
    """
    session = _get_session_or_raise()
    session.log_error(error)


def set_tag(key: str, value: str) -> None:
    """Set a tag on the active session.

    Tags are key-value pairs for organizing and filtering sessions.

    Args:
        key: Tag name
        value: Tag value

    Example:
        armory.set_tag("experiment_type", "ablation")
        armory.set_tag("gpu_type", "A100")
    """
    session = _get_session_or_raise()
    session.set_tag(key, value)


def set_tags(tags: Dict[str, str]) -> None:
    """Set multiple tags on the active session.

    Args:
        tags: Dictionary of tag names to values

    Example:
        armory.set_tags({
            "experiment_type": "ablation",
            "gpu_type": "A100",
            "framework": "pytorch",
        })
    """
    session = _get_session_or_raise()
    session.set_tags(tags)
