"""
Armory SDK - AI-Native MLOps Platform

A Python SDK for logging training metrics, events, and artifacts to Armory.
Designed as a drop-in replacement for Weights & Biases.

Usage:
    import armory

    # Initialize a session
    run = armory.init(
        project="my-project",
        name="experiment-1",
        config={"learning_rate": 1e-4, "epochs": 10}
    )

    # Log metrics
    run.log({"loss": 0.5, "accuracy": 0.9})

    # Log images, tables, histograms
    run.log({
        "sample": armory.Image(img_array),
        "predictions": armory.Table(columns=["id", "pred"], data=[[1, 0.9]]),
        "weights": armory.Histogram(model.weights),
    })

    # Access summary (final metrics)
    run.summary["best_accuracy"] = 0.95

    # Finish the session
    run.finish()

    # Hyperparameter sweeps
    sweep_config = {
        "method": "random",
        "parameters": {"lr": {"min": 1e-5, "max": 1e-2}}
    }
    armory.sweep(sweep_config, function=train, count=10)
"""

# Core run functionality
from armory.run import (
    init,
    Run,
    finish,
    Summary,
    Config,
    MetricDefinition,
)

# Logging configuration
from armory.logger import set_log_level, get_logger

# Callback integrations (optional, may fail if pytorch_lightning is not available)
try:
    from armory.callback import ArmoryCallback
except Exception:
    ArmoryCallback = None

# System metrics
from armory.system import (
    SystemMetrics,
    SystemMonitor,
    collect_system_metrics,
)

# Media types
from armory.media import (
    Image,
    Audio,
    Video,
    Histogram,
    Table,
    Html,
    Object3D,
)

# Sweeps
from armory.sweep import (
    Sweep,
    SweepConfig,
    sweep,
    agent,
)

__version__ = "0.2.0"

__all__ = [
    # Core
    "init",
    "Run",
    "finish",
    "Summary",
    "Config",
    "MetricDefinition",
    # Logging
    "set_log_level",
    "get_logger",
    # Callbacks
    "ArmoryCallback",
    # System
    "SystemMetrics",
    "SystemMonitor",
    "collect_system_metrics",
    # Media types
    "Image",
    "Audio",
    "Video",
    "Histogram",
    "Table",
    "Html",
    "Object3D",
    # Sweeps
    "Sweep",
    "SweepConfig",
    "sweep",
    "agent",
]


# Convenience: make armory.log work like wandb.log
def log(data, step=None, commit=True):
    """Log metrics to the current run"""
    if Run._current:
        Run._current.log(data, step=step, commit=commit)
    else:
        raise RuntimeError("No active run. Call armory.init() first.")


def define_metric(name, step_metric=None, summary="last", goal=None):
    """Define how a metric should be tracked"""
    if Run._current:
        Run._current.define_metric(name, step_metric, summary, goal)


def watch(models, log="gradients", log_freq=1000, log_graph=False):
    """Watch a model for gradient logging"""
    if Run._current:
        Run._current.watch(models, log, log_freq, log_graph)


def alert(title, text, level="INFO", wait_duration=0):
    """Send an alert"""
    if Run._current:
        Run._current.alert(title, text, level, wait_duration)


def save(path, policy="live"):
    """Save a file to the run"""
    if Run._current:
        Run._current.save(path, policy)


def log_artifact(artifact, name=None, type="dataset"):
    """Log an artifact"""
    if Run._current:
        Run._current.log_artifact(artifact, name, type)


# Add to __all__
__all__.extend([
    "log",
    "define_metric",
    "watch",
    "alert",
    "save",
    "log_artifact",
])
