import typing
from typing import TYPE_CHECKING, Any, TypeVar

from zyplib.utils.ensure import ensure_scalar
from zyplib.utils.kv import flatten_dict, prefixed_dict

if TYPE_CHECKING:
    from lightning_fabric.fabric import Fabric

ScalarType = TypeVar('ScalarType', int, float, str, bool)


class LogScalarManager:
    def __init__(self, fabric: 'Fabric'):
        self.fabric = fabric
        self._steps: dict[str, int] = {}

    @property
    def current_steps(self) -> dict[str, int]:
        """返回当前所有 scope 的步数字典"""
        return self._steps

    @current_steps.setter
    def current_steps(self, steps: typing.Mapping[str, int]) -> None:
        self._steps = dict(steps)

    @staticmethod
    def scalar_dict(metrics: typing.Mapping[str, Any]) -> dict[str, Any]:
        """确保 metrics 中的所有值都是标量; 如果无法转换为标量，则直接丢弃"""
        scalar_metrics: dict[str, Any] = {}
        for key, value in metrics.items():
            scalar_value = ensure_scalar(value)
            if scalar_value is not None:
                scalar_metrics[key] = scalar_value
            elif isinstance(value, dict):
                # 如果 value 是 dict，那么就展平它
                flattend_dict = flatten_dict(value)
                scalar_metrics.update(prefixed_dict(f'{key}.', flattend_dict))
            elif value is not None:
                # 这里的 value 就不一定是 scalar 了, 可能是 tensor, 只能祈祷 fabric 能正确处理...
                scalar_metrics[key] = value
        return scalar_metrics

    def _get_step(self, scope: str) -> int:
        if scope not in self._steps:
            self._steps[scope] = 0
        return self._steps[scope]

    def _increment_step(self, scope: str) -> None:
        if scope not in self._steps:
            self._steps[scope] = 0
        self._steps[scope] += 1

    def log_scalars(
        self,
        metrics: typing.Mapping[str, Any],
        *,
        prefix: str = '',
        step: int | None = None,
    ) -> None:
        """记录标量指标到 fabric loggers

        Parameters
        ----------
        - `metrics` : `dict[str, Any]`
            - 要记录的指标字典
        - `prefix` : `str`, optional
            - 指标前缀，如 'train/batch/'，默认为 ''
        - `step` : `int | None`, optional
            - 步数，如果为 None 则自动按 prefix 分 scope 递增
        """
        if prefix:
            metrics = prefixed_dict(prefix, metrics)
        scalar_metrics = self.scalar_dict(metrics)
        if step is None:
            _step = self._get_step(prefix)
        else:
            _step = step
        self.fabric.log_dict(scalar_metrics, step=_step)
        if step is None:
            self._increment_step(prefix)
        else:
            self._steps[prefix] = step + 1
