"""训练信息打印工具

提供独立的训练信息打印逻辑，支持输出到控制台和文件。
纯打印工具，不包含任何 hook 或计时逻辑。
"""

import typing

from pathlib import Path
from typing import Any, Dict, Optional

from zyplib.utils.ensure import ensure_scalar
from zyplib.utils.fs import ensure_dir
from zyplib.utils.kv import flatten_dict
from zyplib.utils.print import print_info


class TrainInfoPrinter:
    def __init__(self, out_file: Optional[str | Path] = None, max_cols: int = 5):
        """训练信息打印器

        提供格式化打印训练信息的功能，支持输出到控制台和文件。
        这是一个纯打印工具类，不包含任何 hook 或计时逻辑。

        Parameters
        ----------
        - `out_file` : `str | Path | None`, optional
            - 输出文件路径，如果提供则会将日志写入文件，默认为 None
        - `max_cols` : `int`, optional
            - 每行最多显示的指标数量，默认为 5

        Examples
        --------
        ```python
        # 基本用法
        printer = TrainInfoPrinter(out_file='./logs/train.log')

        # 打印标题
        printer.print_header('开始训练')

        # 打印指标
        printer.print_metrics({'loss': 0.5, 'acc': 0.9}, prefix='Train')

        # 打印分隔线
        printer.print_separator('Epoch 1/100')
        ```
        """
        self.max_cols = max_cols
        self.out = Path(out_file) if out_file else None

        self._fhandler = None
        if self.out:
            ensure_dir(self.out.parent)
            self._fhandler = open(self.out, 'a')

    def log(self, message: str, end: str = '\n'):
        """输出日志，同时写入控制台和文件"""
        print_info(message, end=end)
        if self._fhandler:
            self._fhandler.write(message + end)

    def __del__(self):
        self.close()

    def close(self):
        """关闭文件句柄"""
        if self._fhandler:
            self._fhandler.close()
            self._fhandler = None

    def flush(self):
        """刷新文件缓冲区"""
        if self._fhandler:
            self._fhandler.flush()

    def print_header(self, text: str, width: int = 120, char: str = '='):
        """打印居中的标题行"""
        self.log(text.center(width, char))

    def print_separator(self, text: str, width: int = 50, char: str = '-'):
        """打印居中的分隔线"""
        self.log(text.center(width, char))

    def print_metrics(self, metrics: typing.Mapping[str, Any], prefix: str = ''):
        """打印指标摘要

        Parameters
        ----------
        - `metrics` : `Dict[str, Any]`
            - 指标字典
        - `prefix` : `str`, optional
            - 前缀标签（如 'Train', 'Val', 'Test'），默认为 ''
        """
        if prefix:
            self.log(f'{prefix:6}|', end='')
        self._print_metrics_inline(metrics)
        self.flush()

    def _print_metrics_inline(self, metrics: typing.Mapping[str, float]):
        """打印指标（内联格式）"""
        i = 0
        metrics = flatten_dict(metrics)
        for key, value in metrics.items():
            if i > 0 and i % self.max_cols == 0:
                self.log('\n' + ' ' * 7, end='')
            scalar = ensure_scalar(value)
            if scalar is not None:
                if 1e-3 < abs(scalar) < 1e5:
                    self.log(f'\t{key}: {scalar:.4f}', end='')
                else:
                    self.log(f'\t{key}: {scalar:.3e}', end='')
            i += 1
        self.log('')  # 换行
