"""
序列化工具模块，支持 pickle 和 joblib 两种方式

传统的 pickle/joblib 只保存一个二进制文件，无可读性
本模块会同时捆绑一个 meta 文件保存基本元信息

pickle 格式：
- filename.pkl
- filename.pkl.meta.yml (自动设置只读)

joblib 格式：
- filename.joblib
- filename.joblib.meta.yml (自动设置只读)
"""

import typing

import os
import pickle
import stat
from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING, Literal

import yaml

from zyplib import import_lazy
from zyplib.utils.print import type_summary

if TYPE_CHECKING:
    import joblib as jbl
else:
    jbl = import_lazy('joblib')

__all__ = [
    'save_pkl',
    'load_pkl',
    'view_pkl_meta',
    'save_joblib',
    'load_joblib',
    'view_joblib_meta',
]

STR_TRUNCAT_LEN = 128


def pkl_path(path: str | Path, type: Literal['pkl', 'meta'] = 'pkl') -> str:
    """获取 pickle 文件路径

    Parameters
    ----------
    - `path` : `str | Path`
        - 基础路径
    - `type` : `Literal['pkl', 'meta']`, optional
        - 路径类型, by default 'pkl'

    Returns
    ----------
    - `str`
        - 完整文件路径
    """
    return _get_file_paths(path, 'pkl', type)


def joblib_path(path: str | Path, type: Literal['joblib', 'meta'] = 'joblib') -> str:
    """获取 joblib 文件路径

    Parameters
    ----------
    - `path` : `str | Path`
        - 基础路径
    - `type` : `Literal['joblib', 'meta']`, optional
        - 路径类型, by default 'joblib'

    Returns
    ----------
    - `str`
        - 完整文件路径
    """
    return _get_file_paths(path, 'joblib', type)


def _get_file_paths(
    path: str | Path,
    format_type: Literal['pkl', 'joblib'],
    file_type: Literal['pkl', 'meta', 'joblib'] | None = None,
) -> str:
    """内部函数：获取文件路径

    Parameters
    ----------
    - `path` : `str | Path`
        - 基础路径
    - `format_type` : `Literal['pkl', 'joblib']`
        - 序列化格式类型
    - `file_type` : `Literal['pkl', 'meta', 'joblib']`, optional
        - 文件类型, by default None (使用 format_type)

    Returns
    ----------
    - `str`
        - 完整文件路径
    """
    path = Path(path)

    if file_type is None:
        file_type = format_type

    if format_type == 'pkl':
        data_path = path.with_suffix('.pkl')
        if file_type == 'pkl':
            return str(data_path)
        elif file_type == 'meta':
            return str(data_path.with_suffix('.pkl.meta.yml'))
    elif format_type == 'joblib':
        data_path = path.with_suffix('.joblib')
        if file_type == 'joblib':
            return str(data_path)
        elif file_type == 'meta':
            return str(data_path.with_suffix('.joblib.meta.yml'))

    raise ValueError(f'Unknown format_type: {format_type} or file_type: {file_type}')


def _generate_meta_info(values: typing.Mapping, data_file_name: str) -> dict:
    """内部函数：生成元信息字典

    Parameters
    ----------
    - `values` : `dict`
        - 要保存的对象字典
    - `data_file_name` : `str`
        - 数据文件名

    Returns
    ----------
    - `dict`
        - 元信息字典
    """
    meta_info = {
        'saved_at': datetime.now().isoformat(),
        'data_file': data_file_name,
        'total_objects': len(values),
        'objects': {},
    }

    # 为每个对象生成智能摘要
    for key, value in values.items():
        try:
            if isinstance(value, (int, float, bool)):
                # 简单类型直接保存原值
                meta_info['objects'][key] = value
            elif isinstance(value, str):
                # 字符串根据长度处理
                if len(value) <= STR_TRUNCAT_LEN:
                    meta_info['objects'][key] = value
                else:
                    meta_info['objects'][key] = {
                        # 'type': 'str',
                        'length': len(value),
                        'preview': value[:STR_TRUNCAT_LEN] + '...',
                        # 'truncated': True,
                    }
            else:
                # 复杂类型使用 type_summary
                meta_info['objects'][key] = type_summary(value, inspect_mode=True)
        except Exception as e:
            # 如果某个对象的摘要生成失败，记录错误信息
            meta_info['objects'][key] = {
                'type': type(value).__name__,
                'summary': f'<摘要生成失败: {str(e)}>',
            }

    return meta_info


def _save_meta_file(meta_path: Path, meta_info: typing.Mapping) -> None:
    """内部函数：保存元信息文件

    Parameters
    ----------
    - `meta_path` : `Path`
        - 元信息文件路径
    - `meta_info` : `dict`
        - 元信息字典
    """
    # 如果文件存在且为只读，先取消只读属性
    if meta_path.exists():
        try:
            os.chmod(meta_path, stat.S_IWRITE | stat.S_IREAD)
        except Exception:
            pass  # 忽略权限问题，可能是文件系统不支持

    try:
        with open(meta_path, 'w', encoding='utf-8') as f:
            yaml.dump(
                meta_info, f, default_flow_style=False, allow_unicode=True, indent=2
            )
    except Exception as e:
        raise RuntimeError(f'保存元信息文件失败: {e}') from e

    # 设置元信息文件为只读
    try:
        os.chmod(meta_path, stat.S_IRUSR | stat.S_IRGRP | stat.S_IROTH)
    except Exception:
        pass  # 忽略权限设置失败，可能是文件系统不支持


def save_pkl(path: str | Path, **values):
    """
    保存对象到 pickle 文件，并生成可读的元信息文件

    保存文件到 pickle，同时生成一个 YAML 格式的元信息文件，方便查看文件内容概要。
    对于不同类型的对象采用智能摘要策略：
    - 简单类型 (int, float, bool)：直接保存原值
    - 字符串：≤30字符保存原值，>30字符截断并标记
    - 复杂类型：使用 type_summary 生成摘要

    Parameters
    ----------
    - `path` : `str | Path`
        - 保存路径，不需要包含扩展名
    - `**values` :
        - 要保存的命名对象，至少需要一个对象

    Raises
    ----------
    - `ValueError`
        - 如果没有提供任何对象保存
    - `PermissionError`
        - 如果无法写入文件

    Examples
    --------
    ```python
    import numpy as np
    data = np.random.rand(100, 10)
    model_weights = {'layer1': np.random.rand(10, 5), 'layer2': np.random.rand(5, 1)}

    save_pkl('experiment_results',
             data=data,
             weights=model_weights,
             config={'lr': 0.001, 'epochs': 100})
    # 将生成：
    # experiment_results.pkl (二进制数据)
    # experiment_results.pkl.meta.yml (可读元信息，只读)
    ```
    """
    if not values:
        raise ValueError('至少需要提供一个对象来保存')

    path = Path(path)

    # 确保目录存在
    path.parent.mkdir(parents=True, exist_ok=True)

    # 保存 pickle 文件
    pkl_path = path.with_suffix('.pkl')
    try:
        with open(pkl_path, 'wb') as f:
            pickle.dump(values, f)
    except Exception as e:
        raise RuntimeError(f'保存 pickle 文件失败: {e}') from e

    # 生成并保存元信息
    meta_info = _generate_meta_info(values, pkl_path.name)
    meta_path = Path(pkl_path.with_suffix('.pkl.meta.yml'))
    _save_meta_file(meta_path, meta_info)

    return pkl_path, meta_path


def load_pkl(path: str | Path) -> dict:
    """
    从 pickle 文件加载对象

    Parameters
    ----------
    - `path` : `str | Path`
        - 文件路径，可以包含或不包含 .pkl 扩展名

    Returns
    ----------
    - `dict`
        - 包含所有保存对象的字典

    Examples
    --------
    ```python
    # 加载之前保存的数据
    loaded = load_pkl('experiment_results')
    data = loaded['data']
    weights = loaded['weights']
    config = loaded['config']
    ```
    """
    path = Path(path)
    if str(path).endswith('.pkl.meta.yml'):
        pkl_path = Path(str(path).replace('.pkl.meta.yml', '.pkl'))
    else:
        pkl_path = path.with_suffix('.pkl')

    if not pkl_path.exists():
        raise FileNotFoundError(f'Pickle 文件不存在: {pkl_path}')

    with open(pkl_path, 'rb') as f:
        values = pickle.load(f)

    # print(f'✓ 已加载: {pkl_path}')
    # print(f'✓ 共加载 {len(values)} 个对象: {list(values.keys())}')

    return values


def view_pkl_meta(path: str | Path) -> dict:
    """
    查看 pickle 文件的元信息，无需加载实际数据

    Parameters
    ----------
    - `path` : `str | Path`
        - 文件路径，可以包含或不包含 .meta.yml 扩展名

    Returns
    ----------
    - `dict`
        - 元信息字典

    Examples
    --------
    ```python
    # 快速查看文件内容摘要
    meta = view_pkl_meta('experiment_results')
    print(f"保存时间: {meta['saved_at']}")
    print(f"对象列表: {list(meta['objects'].keys())}")
    ```
    """
    path = Path(path)
    if path.suffix == '.pkl':
        meta_path = path.with_suffix('.pkl.meta.yml')
    elif path.name.endswith('.meta.yml'):
        meta_path = path
    else:
        # 假设传入的是基础名称，先尝试 .pkl.meta.yml
        meta_path = path.with_suffix('.pkl.meta.yml')
        # 如果不存在，再尝试旧的 .meta.yml 格式（向后兼容）
        if not meta_path.exists():
            meta_path = path.with_suffix('.meta.yml')

    if not meta_path.exists():
        raise FileNotFoundError(f'元信息文件不存在: {meta_path}')

    with open(meta_path, 'r', encoding='utf-8') as f:
        meta_info = yaml.safe_load(f)

    return meta_info


def save_joblib(path: str | Path, **values):
    """
    使用 joblib 保存对象到文件，并生成可读的元信息文件

    joblib 在处理 NumPy 数组时通常比 pickle 更高效，特别适合科学计算场景。
    同时生成一个 YAML 格式的元信息文件，方便查看文件内容概要。
    对于不同类型的对象采用智能摘要策略：
    - 简单类型 (int, float, bool)：直接保存原值
    - 字符串：≤30字符保存原值，>30字符截断并标记
    - 复杂类型：使用 type_summary 生成摘要

    Parameters
    ----------
    - `path` : `str | Path`
        - 保存路径，不需要包含扩展名
    - `**values` :
        - 要保存的命名对象，至少需要一个对象

    Raises
    ----------
    - `ValueError`
        - 如果没有提供任何对象保存
    - `PermissionError`
        - 如果无法写入文件
    - `ImportError`
        - 如果 joblib 没有安装

    Examples
    --------
    ```python
    import numpy as np
    data = np.random.rand(100, 10)
    model_weights = {'layer1': np.random.rand(10, 5), 'layer2': np.random.rand(5, 1)}

    save_joblib('experiment_results',
                data=data,
                weights=model_weights,
                config={'lr': 0.001, 'epochs': 100})
    # 将生成：
    # experiment_results.joblib (二进制数据)
    # experiment_results.joblib.meta.yml (可读元信息，只读)
    ```
    """
    if not values:
        raise ValueError('至少需要提供一个对象来保存')

    path = Path(path)

    # 确保目录存在
    path.parent.mkdir(parents=True, exist_ok=True)

    # 保存 joblib 文件
    joblib_path = path.with_suffix('.joblib')
    try:
        jbl.dump(values, joblib_path)
    except Exception as e:
        raise RuntimeError(f'保存 joblib 文件失败: {e}') from e

    # 生成并保存元信息
    meta_info = _generate_meta_info(values, joblib_path.name)
    meta_path = Path(joblib_path.with_suffix('.joblib.meta.yml'))
    _save_meta_file(meta_path, meta_info)

    return joblib_path, meta_path


def load_joblib(path: str | Path) -> dict:
    """
    从 joblib 文件加载对象

    Parameters
    ----------
    - `path` : `str | Path`
        - 文件路径，可以包含或不包含 .joblib 扩展名

    Returns
    ----------
    - `dict`
        - 包含所有保存对象的字典

    Raises
    ----------
    - `FileNotFoundError`
        - 如果文件不存在
    - `ImportError`
        - 如果 joblib 没有安装

    Examples
    --------
    ```python
    # 加载之前保存的数据
    loaded = load_joblib('experiment_results')
    data = loaded['data']
    weights = loaded['weights']
    config = loaded['config']
    ```
    """
    path = Path(path)
    if str(path).endswith('.joblib.meta.yml'):
        joblib_path = Path(str(path).replace('.joblib.meta.yml', '.joblib'))
    else:
        joblib_path = path.with_suffix('.joblib')

    if not joblib_path.exists():
        raise FileNotFoundError(f'Joblib 文件不存在: {joblib_path}')

    values = jbl.load(joblib_path)

    # print(f'✓ 已加载: {joblib_path}')
    # print(f'✓ 共加载 {len(values)} 个对象: {list(values.keys())}')

    return values


def view_joblib_meta(path: str | Path) -> dict:
    """
    查看 joblib 文件的元信息，无需加载实际数据

    Parameters
    ----------
    - `path` : `str | Path`
        - 文件路径，可以包含或不包含 .meta.yml 扩展名

    Returns
    ----------
    - `dict`
        - 元信息字典

    Raises
    ----------
    - `FileNotFoundError`
        - 如果元信息文件不存在

    Examples
    --------
    ```python
    # 快速查看文件内容摘要
    meta = view_joblib_meta('experiment_results')
    print(f"保存时间: {meta['saved_at']}")
    print(f"对象列表: {list(meta['objects'].keys())}")
    ```
    """
    path = Path(path)
    if path.suffix == '.joblib':
        meta_path = path.with_suffix('.joblib.meta.yml')
    elif path.name.endswith('.meta.yml'):
        meta_path = path
    else:
        # 假设传入的是基础名称，先尝试 .joblib.meta.yml
        meta_path = path.with_suffix('.joblib.meta.yml')

    if not meta_path.exists():
        raise FileNotFoundError(f'元信息文件不存在: {meta_path}')

    with open(meta_path, 'r', encoding='utf-8') as f:
        meta_info = yaml.safe_load(f)

    return meta_info
