"""训练可视化：SwanLab 封装。

配置示例（`train.swanlab`）：

```yaml
train:
  swanlab:
    enabled: true
    project: film-defect
    name: baseline          # 缺省用 work_dir 的目录名
    mode: online            # online(云端) | local(本地看板) | offline(仅存盘) | disabled
    tags: [tiny_conv, center-head]
    log_dir: null           # 缺省 <work_dir>/swanlog
```

未安装 swanlab、或 `enabled: false` 时整体退化为空操作，训练流程不受影响。
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

from utils.logging import get_logger

logger = get_logger("engine.visualizer")

VALID_MODES = ("online", "local", "offline", "disabled")


class SwanLabLogger:
    """把标量指标写到 SwanLab；不可用时静默降级。

    Args:
        enabled: 关掉后所有方法都是空操作。
        work_dir: 用于推导默认的实验名与 `log_dir`。
        config: 记录到实验里的超参（一般直接传完整 yaml 配置）。
        其余参数透传给 `swanlab.init`。
    """

    def __init__(
        self,
        enabled: bool = True,
        work_dir: str | Path = "runs",
        config: dict | None = None,
        project: str = "film-defect",
        name: str | None = None,
        mode: str = "online",
        log_dir: str | Path | None = None,
        **init_kwargs: Any,
    ) -> None:
        work_dir = Path(work_dir)
        self.run = None
        self.enabled = bool(enabled) and mode != "disabled"
        if not self.enabled:
            logger.info("SwanLab 可视化未启用")
            return

        if mode not in VALID_MODES:
            raise ValueError(f"未知 swanlab mode: {mode}，可选 {VALID_MODES}")

        try:
            import swanlab
        except ImportError:
            logger.warning("未安装 swanlab，跳过可视化（uv sync 后可用；或设 train.swanlab.enabled=false）")
            self.enabled = False
            return

        try:
            self.run = swanlab.init(
                project=project,
                name=name or work_dir.name,
                mode=mode,
                config=config or {},
                log_dir=str(log_dir or work_dir / "swanlog"),
                **init_kwargs,
            )
            self._swanlab = swanlab
            logger.info("SwanLab 已启动（project=%s, mode=%s）", project, mode)
        except Exception as exc:  # 联网失败 / 未登录不应中断训练
            logger.warning("SwanLab 初始化失败，训练继续但不记录可视化: %s", exc)
            self.enabled = False

    def log(self, metrics: dict[str, float], step: int | None = None, prefix: str = "") -> None:
        if not self.enabled:
            return
        payload = {f"{prefix}{key}": float(value) for key, value in metrics.items()}
        try:
            self._swanlab.log(payload, step=step)
        except Exception as exc:  # 记录失败同样不中断训练
            logger.warning("SwanLab 记录失败: %s", exc)
            self.enabled = False

    def finish(self) -> None:
        if not self.enabled:
            return
        try:
            self._swanlab.finish()
        except Exception as exc:
            logger.warning("SwanLab 收尾失败: %s", exc)