"""
文件系统相关工具
"""

import typing

import inspect
import json
import os
import shutil
from pathlib import Path
from typing import Any, Callable, Literal, Union

import yaml

from .print import print_debug

FilePath = Union[Path, str]


def expand_user(
    path: FilePath, return_type: Literal['same', 'str', 'path'] = 'same'
) -> Union[str, Path]:
    """将用户主目录展开"""
    result_path = Path(path).expanduser()
    if return_type == 'same':
        return result_path if isinstance(path, Path) else str(result_path)
    elif return_type == 'str':
        return str(result_path)
    elif return_type == 'path':
        return result_path


def expand_relative(
    path: FilePath,
    based_on: FilePath | None = None,
    return_type: Literal['same', 'str', 'path'] = 'same',
) -> Union[str, Path]:
    """完全基于调用文件位置而非工作目录的相对路径

    Parameters
    ----------
    - `path` : `FilePath`
        - 路径
    - `based_on` : `FilePath | None`, optional
        - 基准路径, 会自动识别 based_on 为文件还是目录，并相应地设置基准目录
        - 如果为 None 则通过 `inspect.stack()` 自动识别调用文件位置
    - `return_type` : `Literal['same', 'str', 'path']`, optional
        - 返回类型; 如果为 'same' 则返回 Path 对象, 如果为 'str' 则返回字符串, 如果为 'path' 则返回 Path 对象

    Returns
    ----------
    - `Path`
        - 返回的路径

    Example
    --------
    ```py
    expand_relative('../data/')  # 在 /home/ 下运行的 /home/user/a/test.py 文件
    # /home/user/data/
    expand_relative('../data/', based_on=__file__)
    ```
    """
    if based_on is None:
        caller_file = inspect.stack()[1].filename
        file_dir = Path(caller_file).parent
    else:
        base_path = Path(based_on)
        if base_path.is_file():
            file_dir = base_path.parent
        else:
            file_dir = base_path
    result_path = (file_dir / path).resolve()
    if return_type == 'same':
        return result_path if isinstance(path, Path) else str(result_path)
    elif return_type == 'str':
        return str(result_path)
    elif return_type == 'path':
        return result_path


def ensure_dir(
    path: FilePath,
    if_not_empty: Literal['keep', 'clean', 'error'] = 'keep',
    verbose: bool = False,
) -> Path:
    """检查用于保存日志、Checkpoint 等的目录如果 dir 不存在，则创建

    Parameters
    ----------
    - `path` : `str`
        - 目录
    - `if_not_empty` : `Literal['keep', 'clean', 'error']`, optional
        - 如果目录不为空，则执行的操作
        - `"keep"` | `"clean"` | `"error"`
    - `verbose` : `bool`, optional
        - 是否打印日志，默认为 False
    """
    # 1. 将路径转换为绝对路径，避免相对路径带来的混淆
    p = Path(path).resolve()

    # 2. 区分文件和目录
    directory = p.parent if p.suffix else p

    # 3. 原子性地创建目录
    directory.mkdir(parents=True, exist_ok=True)

    # 4. 检查目录是否为空
    is_not_empty = any(directory.iterdir())

    if is_not_empty:
        if if_not_empty == 'keep':
            if verbose:
                print(f'Directory exists and is not empty. Kept as is: {directory}')
        elif if_not_empty == 'clean':
            if verbose:
                print(f'Directory exists and is not empty. Cleaning up: {directory}')
            shutil.rmtree(directory)
            directory.mkdir()  # 此时父目录一定存在，无需 parents=True
        elif if_not_empty == 'error':
            raise FileExistsError(f'Directory is not empty: {directory}')
        else:
            # 这个分支理论上在 Literal 类型约束下很难触发，但作为保障依然保留
            raise ValueError(f"Invalid value for 'if_not_empty': {if_not_empty}")

    return directory


def file_reader(file_path: FilePath, encoding: str = 'utf-8', binary: bool = False):
    mode = 'rb' if binary else 'r'
    return open(file_path, mode, encoding=None if binary else encoding)


def read_file(
    file_path: FilePath,
    encoding: str = 'utf-8',
    binary: bool = False,
    readlines: bool = False,
) -> Union[str, list[str]]:
    with file_reader(file_path, encoding, binary) as f:
        if readlines:
            return f.readlines()
        return f.read()


def file_writer(file_path: FilePath, encoding: str = 'utf-8', binary: bool = False):
    mode = 'wb' if binary else 'w'
    return open(file_path, mode, encoding=None if binary else encoding)


def write_file(
    file_path: FilePath,
    content: Union[str, bytes],
    encoding: str = 'utf-8',
    binary: bool = False,
):
    ensure_dir(os.path.dirname(file_path))
    with file_writer(file_path, encoding, binary) as f:
        f.write(content)


def edit_file_lines(
    file_path: FilePath,
    fn_edit_line: Callable[[str], str],
    fn_pre_filter_line: Callable[[str], bool] | None = None,
    fn_post_filter_line: Callable[[str], bool] | None = None,
    encoding: str = 'utf-8',
):
    """编辑文件的每一行

    Parameters
    ----------
    - `file_path` : `FilePath`
        - 文件路径
    - `fn_edit_line` : `Callable[[str], str]`
        - 编辑每一行的函数
    - `fn_pre_filter_line` : `Callable[[str], bool] | None`, optional
        - 预过滤函数, 如果返回 False 则跳过该行, 默认为 None
    - `fn_post_filter_line` : `Callable[[str], bool] | None`, optional
        - 后过滤函数, 如果返回 False 则不写入该行, 默认为 None
    - `encoding` : `str`, optional
        - 编码方式, 默认为 'utf-8'
    """
    lines = read_file(file_path, encoding=encoding, readlines=True)
    new_lines = []
    for line in lines:
        if fn_pre_filter_line and not fn_pre_filter_line(line):
            continue
        new_line = fn_edit_line(line)
        if fn_post_filter_line and not fn_post_filter_line(new_line):
            continue
        new_lines.append(new_line)
    write_file(file_path, ''.join(new_lines), encoding=encoding)


def read_json(file_path: FilePath, encoding: str = 'utf-8'):
    with open(file_path, 'r', encoding=encoding) as f:
        return json.load(f)


def write_json(
    file_path: FilePath, data: typing.Mapping | typing.Sequence, encoding: str = 'utf-8'
):
    ensure_dir(os.path.dirname(file_path))
    with open(file_path, 'w', encoding=encoding) as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


# # 自定义标签处理函数，将其转换为字符串
# def ignore_unknown_tag(loader, node):
#     return loader.construct_scalar(node)


# # 注册 pathlib.PosixPath 的处理方式为忽略
# yaml.SafeLoader.add_constructor(
#     'tag:yaml.org,2002:python/object/apply:pathlib.PosixPath', ignore_unknown_tag
# )


def read_yaml(file_path: FilePath, throw_error: bool = True):
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            return yaml.safe_load(f)
    except yaml.YAMLError as e:
        if throw_error:
            raise e
        return {}


# 定义字符串呈现器（处理多行字符串）
def str_presenter(dumper, data):
    if len(data.splitlines()) > 1:  # check for multiline string
        return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|')
    return dumper.represent_scalar('tag:yaml.org,2002:str', data)


# 注册 str 类型的全局呈现器（只需要注册一次）
yaml.add_representer(str, str_presenter)


def write_yaml(file_path: FilePath, data: typing.Mapping):
    # 确保导出时只使用基本的数据类型
    def dump_for_yaml(obj: Any) -> Any:
        if isinstance(obj, (int, float, bool, str)):
            return obj
        elif isinstance(obj, dict):
            return {k: dump_for_yaml(v) for k, v in obj.items()}
        elif isinstance(obj, (list, tuple, set)):
            return [dump_for_yaml(i) for i in obj]
        elif obj is None:
            return None
        elif hasattr(obj, 'to_dict'):  # 支持自定义对象的序列化
            return dump_for_yaml(obj.to_dict())
        return str(obj)

    clean_data = dump_for_yaml(data)  # 清理自定义标签

    ensure_dir(os.path.dirname(file_path))  # 确保目录存在
    with open(file_path, 'w', encoding='utf-8') as f:
        yaml.dump(clean_data, f, default_flow_style=False, allow_unicode=True)
