"""检查点保存管理器

提供独立的检查点保存逻辑，支持 Top-K 策略、索引管理等功能。
可以独立使用，也可以与 Callback 组合使用。
"""

import typing

import json
import operator
import os
from pathlib import Path
from typing import Any, Callable, Dict, Literal, 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_debug, print_warn, warning


def _smart_format_score(score: float) -> str:
    """智能格式化 score，极端值使用科学计数法

    对于正常范围的数值 (1e-4 ~ 1e4) 使用定点格式 .4f
    对于极端值使用科学计数法 .4e 保持精度

    Parameters
    ----------
    - `score` : `float`
        - 要格式化的分数

    Returns
    -------
    - `str`
        - 格式化后的字符串
    """
    if score == 0:
        return '0.0000'
    abs_score = abs(score)
    if 1e-4 <= abs_score < 1e4:
        return f'{score:.4f}'
    else:
        # 科学计数法，保留4位有效数字
        return f'{score:.4e}'


class TopKCheckpointManager:
    MODEL_FNAME_SEP = '#'

    def __init__(
        self,
        dir: str,
        keep_best_k: int = 1,
        monitor: str = 'loss',
        less_is_better: bool = True,
        ckpt_ext: str = '.ckpt',
        if_not_empty: Literal['ask', 'clean', 'error', 'ignore'] = 'ask',
        save_fn: Optional[Callable[[str], None]] = None,
    ):
        """检查点保存管理器

        提供独立的检查点保存逻辑，支持 Top-K 策略、索引管理等功能。

        Parameters
        ----------
        - `dir` : `str`
            - 保存检查点的目录
        - `keep_best_k` : `int`, optional
            - 保留表现最好的 K 个模型，默认为 1
            - Top-K 策略：只要比当前保存的最差模型好，或还没存满 K 个，就会保存
        - `monitor` : `str`, optional
            - 用于确定最佳模型的指标名称，默认为 'loss'
        - `less_is_better` : `bool`, optional
            - 是否为越小越优，默认为 True
        - `ckpt_ext` : `str`, optional
            - 检查点文件的扩展名，默认为 '.ckpt'
        - `if_not_empty` : `Literal['ask', 'clean', 'error', 'ignore']`, optional
            - 当检查点目录不为空时的处理方式，默认为 'ask'
            - 'ask': 交互式询问用户（注意：非交互环境会抛出异常）
            - 'clean': 直接清空目录
            - 'error': 抛出异常
            - 'ignore': 忽略，继续使用现有目录
        - `save_fn` : `Callable[[str], None]`, optional
            - 自定义保存函数，接收检查点路径作为参数
            - 如果不提供，需要在调用 `update` 时传入 `save_fn`

        Examples
        --------
        ```python
        # 基本用法
        manager = CheckpointTopKSaving(
            dir='./checkpoints',
            keep_best_k=3,
            monitor='val_loss',
            less_is_better=True,
        )

        # 在训练循环中
        for epoch in range(100):
            # ... 训练代码 ...
            val_loss = validate()

            # 更新检查点
            manager.update(
                score=val_loss,
                epoch=epoch,
                save_fn=lambda path: torch.save(model.state_dict(), path)
            )

        # 获取最佳检查点
        best_ckpt = manager.best_checkpoint()
        ```
        """
        self._init_dir(dir, if_not_empty)

        self.dir = Path(dir)
        self.keep_best_k = keep_best_k if keep_best_k and keep_best_k > 0 else 1
        self.monitor = monitor
        self.less_is_better = less_is_better
        self.ckpt_ext = ckpt_ext
        self.save_fn = save_fn

        # 比较函数：less_is_better 时使用 lt，否则使用 gt
        self.score_cmp = operator.lt if less_is_better else operator.gt

        # 记录历史最佳分数
        self.best_score = float('inf') if less_is_better else float('-inf')

        # 检查点索引文件
        self._ckpt_index_path = self.dir / '.ckpt_index.json'
        self._ckpt_index: Dict[str, dict] = {}  # filename -> {epoch, monitor, score}

    def _init_dir(
        self, dir: str, if_not_empty: Literal['ask', 'clean', 'error', 'ignore']
    ):
        """初始化检查点目录"""
        import os
        import sys

        ensure_dir(dir)

        # 检查空文件夹
        if os.listdir(dir):
            if if_not_empty == 'ask':
                # 检测是否为交互式环境
                if not sys.stdin.isatty():
                    raise RuntimeError(
                        f'检查点文件夹 {dir} 不为空，且当前为非交互式环境。'
                        f'请使用 if_not_empty="clean"|"error"|"ignore" 参数指定处理方式。'
                    )
                warning(f'检查点文件夹 {dir} 不为空, 可能存在旧的检查点文件')
                user = input('按 y 清空文件夹后继续训练, 否则退出: (y/n)')
                if user.lower().strip() != 'y':
                    exit()
                else:
                    ensure_dir(dir, if_not_empty='clean')
            elif if_not_empty == 'clean':
                warning(f'检查点文件夹 {dir} 不为空, 正在清空...')
                ensure_dir(dir, if_not_empty='clean')
            elif if_not_empty == 'error':
                raise RuntimeError(f'检查点文件夹 {dir} 不为空!')
            # 'ignore' 时什么都不做

    def _process_extra_info(self, extra_info: typing.Mapping[str, Any]) -> Dict[str, Any]:
        """处理额外信息，展平字典结构"""
        info = flatten_dict(extra_info)
        result = {}
        for k, v in info.items():
            result[k] = ensure_scalar(v)
            if not isinstance(result[k], (int, float, complex)):
                result[k] = repr(result[k])
        return result

    def update(
        self,
        score: float,
        epoch: int,
        save_fn: Optional[Callable[[str], None]] = None,
        extra_info: Optional[typing.Mapping[str, Any]] = None,
    ) -> Optional[Path]:
        """更新检查点

        根据 Top-K 策略判断是否需要保存当前检查点。

        Parameters
        ----------
        - `score` : `float`
            - 当前的评估分数
        - `epoch` : `int`
            - 当前的 epoch 数
        - `save_fn` : `Callable[[str], None]`, optional
            - 自定义保存函数，接收检查点路径作为参数
            - 如果不提供，使用初始化时传入的 save_fn
        - `extra_info` : `Dict[str, Any]`, optional
            - 额外信息，会保存到索引文件中

        Returns
        -------
        - `Optional[Path]`
            - 如果保存了检查点，返回检查点路径；否则返回 None
        """
        raw_score = score
        scalar_score = ensure_scalar(score)
        if not isinstance(scalar_score, (int, float, complex)):
            raise ValueError(f'score 必须是标量数值，但收到 {raw_score}')
        score = typing.cast(float, scalar_score)

        if extra_info:
            extra_info = self._process_extra_info(extra_info)

        # 判断是否需要保存
        if not self._should_save_topk(score):
            return None

        # 确定保存函数
        fn = save_fn or self.save_fn
        if fn is None:
            raise ValueError('必须提供 save_fn 参数或在初始化时设置 save_fn')

        # 更新历史最佳分数
        if self.score_cmp(score, self.best_score):
            self.best_score = score

        # 生成检查点文件名并保存
        ckpt_name = self._format_best_model_fname(
            epoch=epoch, monitor=self.monitor, score=score
        )
        ckpt_path = self.dir / ckpt_name
        fn(str(ckpt_path))

        # 更新索引
        index_entry = {
            'epoch': epoch,
            'monitor': self.monitor,
            'score': score,  # 保存原始 float 值，无精度损失
        }
        if extra_info:
            index_entry['extra'] = extra_info
        self._ckpt_index[ckpt_name] = index_entry
        self._save_ckpt_index()

        print_debug(f'Saved best model at epoch {epoch} with {self.monitor}: {score:.4f}')

        # 剔除超出 Top-K 的检查点
        self._prune_checkpoints()

        return ckpt_path

    def _should_save_topk(self, current_score: float) -> bool:
        """判断当前分数是否应该保存到 Top-K

        真正的 Top-K 策略：
        1. 如果还没存满 K 个，直接保存
        2. 如果已存满 K 个，只要比最差的那个好就保存
        """
        existing_ckpts = list(self.dir.glob(f'best_model*{self.ckpt_ext}'))
        current_count = len(existing_ckpts)

        # 还没存满 K 个，直接保存
        if current_count < self.keep_best_k:
            return True

        # 已存满，需要比较当前分数与 Top-K 中最差的分数
        worst_score = self._get_worst_score_in_topk()
        if worst_score is None:
            return True  # 无法获取最差分数，保守起见保存

        # 当前分数比最差的好，则保存
        return self.score_cmp(current_score, worst_score)

    def _get_worst_score_in_topk(self) -> Optional[float]:
        """获取当前 Top-K 中最差的分数"""
        existing_ckpts = list(self.dir.glob(f'best_model*{self.ckpt_ext}'))
        if not existing_ckpts:
            return None

        scores = []
        for ckpt in existing_ckpts:
            ckpt_name = ckpt.name
            if ckpt_name in self._ckpt_index:
                scores.append(self._ckpt_index[ckpt_name]['score'])
            else:
                # 回退到从文件名解析
                info = self._extract_best_model_info(ckpt_name)
                if info['score'] is not None:
                    scores.append(info['score'])

        if not scores:
            return None

        # 最差的分数：如果 less_is_better，最差是最大的；否则最差是最小的
        return max(scores) if self.less_is_better else min(scores)

    def _prune_checkpoints(self):
        """剔除超出 Top-K 的检查点"""

        def get_score_for_sort(ckpt_path: Path) -> float:
            ckpt_name = ckpt_path.name
            if ckpt_name in self._ckpt_index:
                return self._ckpt_index[ckpt_name]['score']
            # 回退到从文件名解析
            info = self._extract_best_model_info(ckpt_name)
            return info.get(
                'score', float('inf') if self.less_is_better else float('-inf')
            )

        checkpoints = sorted(
            self.dir.glob(f'best_model*{self.ckpt_ext}'),
            key=get_score_for_sort,
            reverse=not self.score_cmp(0, 1),  # True if higher is better
        )

        if len(checkpoints) > self.keep_best_k:
            for checkpoint in checkpoints[self.keep_best_k :]:
                self._remove_from_index(checkpoint.name)
                checkpoint.unlink()

    def _format_best_model_fname(
        self,
        epoch: Optional[int] = None,
        monitor: Optional[str] = None,
        score: Optional[float] = None,
    ) -> str:
        """格式化最佳模型文件名"""
        sep = self.MODEL_FNAME_SEP
        epoch_part = f'{sep}epoch={epoch}' if epoch is not None else ''
        monitor_part = f'{sep}on={monitor}' if monitor else ''
        score_part = (
            f'{sep}score={_smart_format_score(score)}' if score is not None else ''
        )

        return f'best_model{epoch_part}{monitor_part}{score_part}{self.ckpt_ext}'

    def _extract_best_model_info(self, fname: str) -> Dict[str, Any]:
        """从文件名中提取检查点信息"""
        fname = fname.removesuffix(self.ckpt_ext)
        each_parts = fname.split(self.MODEL_FNAME_SEP)
        info: Dict[str, Any] = {'epoch': None, 'monitor': None, 'score': None}
        for each_part in each_parts:
            if each_part.startswith('epoch='):
                info['epoch'] = int(each_part.split('=')[-1])
            elif each_part.startswith('on='):
                info['monitor'] = each_part.split('=')[-1]
            elif each_part.startswith('score='):
                info['score'] = float(each_part.split('=')[-1])
        return info

    # ==================== 索引管理 ====================

    def _save_ckpt_index(self):
        """保存检查点索引到 JSON 文件"""
        # with open(self._ckpt_index_path, 'w', encoding='utf-8') as f:
        #     json.dump(self._ckpt_index, f, indent=2, ensure_ascii=False)
        temp_path = self._ckpt_index_path.with_suffix('.tmp')
        with open(temp_path, 'w', encoding='utf-8') as f:
            json.dump(self._ckpt_index, f, indent=2, ensure_ascii=False)
        os.replace(temp_path, self._ckpt_index_path)

    def _load_ckpt_index(self) -> Dict[str, dict]:
        """从 JSON 文件加载检查点索引"""
        if not self._ckpt_index_path.exists():
            return {}
        try:
            with open(self._ckpt_index_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except (json.JSONDecodeError, IOError):
            return {}

    def _remove_from_index(self, ckpt_name: str):
        """从索引中移除检查点记录"""
        if ckpt_name in self._ckpt_index:
            del self._ckpt_index[ckpt_name]
            self._save_ckpt_index()

    # ==================== 查询接口 ====================

    def best_checkpoint(self) -> Path:
        """获取最佳检查点路径

        优先从索引文件中查找（精度更高），如果索引不可用则回退到扫描文件名

        Returns
        -------
        - `Path`
            - 最佳检查点的路径

        Raises
        ------
        - `FileNotFoundError`
            - 如果没有找到任何检查点文件
        """
        checkpoints = list(self.dir.glob(f'best_model*{self.ckpt_ext}'))
        if not checkpoints:
            raise FileNotFoundError(
                f'No checkpoints found in {self.dir} with extension {self.ckpt_ext}'
            )

        # 尝试从索引中查找最佳检查点
        index = self._load_ckpt_index()
        if index:
            # 筛选出实际存在的检查点
            existing_ckpts = {ckpt.name for ckpt in checkpoints}
            valid_entries = [
                (name, info) for name, info in index.items() if name in existing_ckpts
            ]

            if valid_entries:
                # 根据 score 排序找到最佳的
                if self.less_is_better:
                    best_name, _ = min(valid_entries, key=lambda x: x[1]['score'])
                else:
                    best_name, _ = max(valid_entries, key=lambda x: x[1]['score'])
                return self.dir / best_name

        # 回退到从文件名解析 score（可能有精度损失）
        print_warn('未找到检查点索引，回退到文件名扫描模式')

        def extract_score(checkpoint: Path) -> float:
            info = self._extract_best_model_info(checkpoint.name)
            return info.get(
                'score', float('inf') if self.less_is_better else float('-inf')
            )

        best_checkpoint = max(
            checkpoints,
            key=lambda x: extract_score(x)
            if self.score_cmp == operator.gt
            else -extract_score(x),
        )
        return best_checkpoint

    def list_checkpoints(self) -> list[Path]:
        """列出所有最佳模型检查点

        Returns
        -------
        - `list[Path]`
            - 所有最佳模型检查点的路径列表，按分数从好到差排序
        """
        checkpoints = list(self.dir.glob(f'best_model*{self.ckpt_ext}'))

        def get_score(ckpt: Path) -> float:
            name = ckpt.name
            if name in self._ckpt_index:
                return self._ckpt_index[name]['score']
            info = self._extract_best_model_info(name)
            return info.get(
                'score', float('inf') if self.less_is_better else float('-inf')
            )

        # 按分数排序（从好到差）
        return sorted(checkpoints, key=get_score, reverse=not self.less_is_better)

    def get_checkpoint_info(self, ckpt_path: Path) -> Optional[Dict[str, Any]]:
        """获取指定检查点的信息

        Parameters
        ----------
        - `ckpt_path` : `Path`
            - 检查点路径

        Returns
        -------
        - `Optional[Dict[str, Any]]`
            - 检查点信息，包含 epoch、monitor、score 等
        """
        ckpt_name = ckpt_path.name
        if ckpt_name in self._ckpt_index:
            return self._ckpt_index[ckpt_name].copy()

        # 回退到从文件名解析
        return self._extract_best_model_info(ckpt_name)
