from __future__ import annotations

from pathlib import Path
from typing import Any


class ExperimentLogger:
    def __init__(self, config: dict[str, Any], mode: str):
        self.config = config
        self.mode = str(mode)
        self.enabled = bool(config.get("logging", {}).get("wandb", {}).get("enabled", False))
        self._wandb = None
        self._run = None

    def start(self) -> None:
        if not self.enabled:
            return
        try:
            import wandb
        except ImportError as exc:
            raise ImportError("wandb logging is enabled, but wandb is not installed") from exc

        wandb_cfg = dict(self.config.get("logging", {}).get("wandb", {}))
        project = str(wandb_cfg.get("project", "atmos_project"))
        run_name = wandb_cfg.get("run_name") or Path(str(self.config.get("output_dir", "run"))).name
        if self.mode != "train":
            run_name = f"{run_name}_{self.mode}"
        self._wandb = wandb
        self._run = wandb.init(
            project=project,
            name=run_name,
            config={k: v for k, v in self.config.items() if not str(k).startswith("_")},
            tags=wandb_cfg.get("tags"),
            notes=wandb_cfg.get("notes"),
            mode=str(wandb_cfg.get("mode", "online")),
        )

    def log(self, metrics: dict[str, Any], step: int | None = None, prefix: str | None = None) -> None:
        if not self.enabled or self._wandb is None:
            return
        payload = {}
        for key, value in metrics.items():
            if isinstance(value, (int, float, str, bool)) or value is None:
                payload[f"{prefix}/{key}" if prefix else str(key)] = value
        self._wandb.log(payload, step=step)

    def log_file(self, path: str | Path, name: str | None = None) -> None:
        if not self.enabled or self._wandb is None:
            return
        artifact = self._wandb.Artifact(name or Path(path).stem, type=f"{self.mode}_result")
        artifact.add_file(str(path))
        self._wandb.log_artifact(artifact)

    def finish(self) -> None:
        if self.enabled and self._wandb is not None:
            self._wandb.finish()