import typing
import os
from typing import Any, Callable, Literal, Optional, TypedDict

from lightning.fabric import Fabric
from lightning.fabric.loggers.csv_logs import CSVLogger
from lightning.fabric.loggers.logger import Logger
from lightning.fabric.loggers.tensorboard import TensorBoardLogger
from typing_extensions import Required

from zyplib.utils.time import TIME_FMT, now

from . import trainer
from .callbacks import (
    FabricCallbackProtocol,
    FabricLogMetrics,
    ParseCSVLogfile,
    PrintTrainInfo,
    SaveCheckpoints,
    SaveHparams,
    SaveLoggersFinally,
    find_callback,
)
from .procedure import (
    FabricTrainingContext,
    train_epoch,
    train_loop,
    training_step,
    val_epoch,
    validation_step,
)
from .trainer import BasicFabricTrainer

basic_trainer = trainer  # 兼容旧的模块命名

__all__ = [
    'default_fabric',
    # callbacks
    'FabricCallbackProtocol',
    'find_callback',
    # procedure
    'FabricTrainingContext',
    'train_epoch',
    'train_loop',
    'training_step',
    'val_epoch',
    'validation_step',
    # trainer
    'BasicFabricTrainer',
]


class EnabledSideEffect(TypedDict, total=False):
    """控制是否启用各种内置功能的配置"""

    csv_logger: Required[bool]  # CSVLogger
    tensorboard_logger: Required[bool]  # TensorBoardLogger
    print_info_cb: Required[bool]  # PrintTrainInfo callback
    save_train_info: Required[bool]  # SaveTrainInfo callback
    save_checkpoints_cb: Required[bool]  # SaveCheckpoints callback
    save_hparams_cb: Required[bool]  # SaveHparams callback
    # write_train_info: Optional[bool]  # WriteTrainInfo callback


EnabledConfig = EnabledSideEffect | dict[str, bool]


# 默认全部启用的配置
DEFAULT_ENABLED: EnabledSideEffect = {
    'csv_logger': False,
    'tensorboard_logger': True,
    'print_info_cb': True,
    'save_train_info': True,
    'save_checkpoints_cb': True,
    'save_hparams_cb': True,
}


def default_fabric(
    root_dir: str = './',
    dir_name: str = 'lightning_logs',
    version: str | None = None,
    save_every_k: Optional[int] = None,
    save_best: bool = True,
    keep_best_k: Optional[int] = 3,
    monitor: str = 'loss',
    monitor_on: Literal['train', 'val'] = 'val',
    less_is_better: bool = True,
    enabled: EnabledConfig | bool | None = None,
    fn_custom_logger: Optional[Callable[[str], typing.Sequence[Logger]]] = None,
    fn_custom_callback: Optional[
        Callable[[str], typing.Sequence[FabricCallbackProtocol]]
    ] = None,
    **fabric_kwargs,
):
    """构建默认的fabric对象

    将 logger 和 checkpoint 保存在 `root_dir/dir_name/version` 目录下
    默认配置为`./lightning_logs/v2024_05_01-12_00_00`

    默认配置如下:
    - Loggers:
        - CSVLogger
        - TensorBoardLogger
    - Callbacks:
        - PrintTrainInfo
        - SaveHparams
        - SaveCheckpoints

    可以通过 `fn_custom_xxx` 参数来添加自定义的 logger 和 callback; 该回调函数传入 now 的时间字符串，需要返回一个 logger 或者 callback 的列表

    Parameters
    ----------
    - `root_dir` : `str`, optional
        - 日志保存的根目录, by default './'
    - `dir_name` : `str`, optional
        - 日志保存的目录名称, by default 'lightning_logs'
    - `version` : `str | None`, optional
        - 日志保存的版本, by default None
        - 默认为 `v{time_str}`
    - `save_every_k` : `Optional[int]`, SaveCheckpoints 的参数
        - 每多少个epoch保存一次模型, by default None
    - `save_best` : `bool`, SaveCheckpoints 的参数
        - _description_, by default True
    - `keep_best_k` : `Optional[int]`, SaveCheckpoints 的参数
        - 保存最好的k个模型, by default 3
    - `monitor` : `str`, SaveCheckpoints 的参数
        - 监控的指标, by default 'loss'
    - `less_is_better` : `bool`, SaveCheckpoints 的参数
        - 是否监控指标越小越好, by default True
    - `enabled` : `dict[str, bool] | bool | None`, optional
        - 控制是否启用各种内置功能的配置，可以部分指定，默认启用所有内置功能。包含以下字段：
        - `csv_logger`: 是否启用 CSVLogger
        - `tensorboard_logger`: 是否启用 TensorBoardLogger
        - `print_info_cb`: 是否启用打印训练信息回调
        - `save_checkpoints_cb`: 是否启用保存检查点回调
        - `save_hparams_cb`: 是否启用保存超参数回调
        - `*`: 通配符, 用于批量设置, 但是会被上面的覆盖

    - `fn_custom_logger` : `(version: string) => Logger[]`, optional
        - 自定义 logger 的函数, by default None
    - `fn_custom_callback`: `(version: string) => FabricCallbackProtocol[]`
        - 自定义 callback 的函数, by default None

    Returns
    ----------
    - `fabric` : `Fabric`
        - Fabric 对象; 注意只返回 fabric，不自动调用 fabric.launch()

    Examples
    ----------

    ```py
    # 禁用所有内置功能, 只开启打印训练信息回调
    >>> fabric = default_fabric(enabled={'*': False, 'print_info_cb': True})
    >>> fabric.launch()
    ```

    """
    # 合并配置
    config = dict(DEFAULT_ENABLED)
    if enabled is not None:
        if isinstance(enabled, bool):
            config = {k: enabled for k in DEFAULT_ENABLED}
        elif isinstance(enabled, dict):
            enabled_dict = typing.cast(dict[str, bool], enabled)
            if '*' in enabled_dict:
                # 先用通配符设置所有值
                default_enabled = enabled_dict['*']
                config = {k: default_enabled for k in DEFAULT_ENABLED}
                # 然后用具体配置覆盖
                for k, v in enabled_dict.items():
                    if k != '*' and k in config:
                        config[k] = v
            else:
                config.update(enabled_dict)

    time_str = now(TIME_FMT.DATE_TIME_UNDERSCORE)
    version = version or f'v{time_str}'
    version = version.replace(r'{time_str}', time_str)
    _loggers: list[Logger] = []
    _callbacks: list[Any] = []

    # Add loggers if enabled
    if config['csv_logger']:
        csv = CSVLogger(root_dir=root_dir, name=dir_name, version=version)
        _loggers.append(csv)
        parser = ParseCSVLogfile(
            os.path.join(csv.log_dir, 'metrics.csv'), to_excel=True, to_yaml=False
        )
        _callbacks.append(parser)
    if config['tensorboard_logger']:
        _loggers.append(
            TensorBoardLogger(root_dir=root_dir, name=dir_name, version=version)
        )

    if fn_custom_logger:
        custom_loggers = fn_custom_logger(version)
        if isinstance(custom_loggers, list):
            _loggers.extend(custom_loggers)
        else:
            raise ValueError(
                f'fn_custom_logger must return a list of loggers, but got {type(custom_loggers)}'
            )

    # Add callbacks if enabled
    if config['print_info_cb']:
        out_dir = None
        if config['save_train_info']:
            out_dir = f'{root_dir}/{dir_name}/{version}'
        _callbacks.append(PrintTrainInfo(out_dir))
    if config['save_hparams_cb']:
        _callbacks.append(SaveHparams(dir=f'{root_dir}/{dir_name}/{version}'))
    if config['save_checkpoints_cb']:
        _callbacks.append(
            SaveCheckpoints(
                dir=f'{root_dir}/{dir_name}/{version}/checkpoints',
                save_every_k=save_every_k,
                save_best=save_best,
                keep_best_k=keep_best_k,
                monitor=monitor,
                monitor_on=monitor_on,
                less_is_better=less_is_better,
            )
        )
    if fn_custom_callback:
        custom_callbacks = fn_custom_callback(version)
        if isinstance(custom_callbacks, list):
            _callbacks.extend(custom_callbacks)
        else:
            raise ValueError(
                f'fn_custom_callback must return a list of callbacks, but got {type(custom_callbacks)}'
            )

    fabric = Fabric(loggers=_loggers, callbacks=_callbacks, **fabric_kwargs)

    # 如果启用了日志记录, 则添加日志记录回调, 用于自动记录 metrics 到日志中
    if len(_loggers) > 0:
        fabric._callbacks.append(FabricLogMetrics(fabric))
        fabric._callbacks.append(SaveLoggersFinally(fabric))

    return fabric
