"""Key-value 相关工具"""

import typing

from functools import singledispatch
from itertools import chain
from pathlib import Path
from typing import Any, Callable, OrderedDict
from typing_extensions import Literal

from zyplib.utils.fs import read_json, read_yaml
from zyplib.utils.print import print_error

CollectKeyPathResult = dict[str, Any] | OrderedDict[str, dict[str, Any]]


def filter_kwargs(
    d: typing.Mapping[str, Any],
    *,
    allowed: typing.Sequence[str] | typing.AbstractSet[str] | None = None,
    prohibited: typing.Sequence[str] | typing.AbstractSet[str] | None = None,
) -> dict[str, Any]:
    """Filters a dictionary to only include allowed keys."""
    if allowed is None and prohibited is None:
        return typing.cast(dict[str, Any], d)
    # return {k: v for k, v in d.items() if k in allowed}
    if allowed is None:
        allowed = set(d.keys())
    if prohibited is None:
        prohibited = set()
    return {k: v for k, v in d.items() if k in allowed and k not in prohibited}


@singledispatch
def collect_key_path(target: Any, key: str) -> CollectKeyPathResult:
    """获取一个字典中所有指定 key 的对象，并返回其压缩后的对象 {路径: value}

    - 输入参数可以为 dict 对象：对对象进行检索
    - 也可以为一个路径：检索路径下所有 yml 和 json 文件内的 key，并返回汇总结果

    ```py
    >>> dic = {
        'loss': 1,
        'avg': { 'loss': 1.2 },
        'Exp': { 'Round0': { 'loss': 2 }, 'Round1': { 'loss': 1 } }
    }
    >>> collect_key_paths(dic, 'loss')
    {'loss': 1, 'avg.loss': 1.2, 'Exp.Round0.loss': 2, 'Exp.Round1.loss': 1}
    ```
    """
    raise TypeError(f'Unsupported type for collect_key_path: {type(target)}')


@collect_key_path.register(dict)
def _(target: dict[str, Any], key: str) -> dict[str, Any]:
    """处理字典类型的输入"""
    result = {}

    def recurse(sub_dict: typing.Mapping, parent_key: str = ''):
        for k, v in sub_dict.items():
            new_key = f'{parent_key}.{k}' if parent_key else k
            if k == key:
                result[new_key] = v
            if isinstance(v, dict):
                recurse(v, new_key)

    recurse(target)
    return result


def _collect_from_files(target: Path, key: str) -> OrderedDict[str, dict[str, Any]]:
    """Helper to process a directory path."""
    path = Path(target)
    result = OrderedDict()

    # 递归查看目录中的所有 YAML/JSON 文件
    all_files = chain(path.rglob('*.yaml'), path.rglob('*.yml'), path.rglob('*.json'))
    for fpath in all_files:
        content: dict = {}
        try:
            if fpath.suffix in ('.yaml', '.yml'):
                content = read_yaml(fpath)
            elif fpath.suffix == '.json':
                content = read_json(fpath)
            else:
                continue

            # Use the registered dict implementation
            key_paths = collect_key_path(content, key)
            if key_paths:
                result[str(fpath)] = key_paths
        except Exception as e:
            print_error(f'Error reading {fpath}: {repr(e)}')
    return result


@collect_key_path.register(str)
def _(target: str, key: str) -> OrderedDict[str, dict[str, Any]]:
    return _collect_from_files(Path(target), key)


@collect_key_path.register(Path)
def _(target: Path, key: str) -> OrderedDict[str, dict[str, Any]]:
    return _collect_from_files(target, key)


def prefixed_dict(prefix: str, d: typing.Mapping[str, Any]) -> dict[str, Any]:
    """给字典的每个键值对添加前缀

    Example
    --------
    ```py
    >>> dic = torch_metric.compute()
    >>> prefixed_dict('val_', dic)
    >>> {'val_loss': 0.1, 'val_acc': 0.9}
    ```


    Parameters
    ----------
    - `prefix` : `str`
        - 前缀
    - `d` : `dict`
        - 需要添加前缀的字典

    Returns
    ----------
    - `dict`
        - 添加前缀后的字典
    """
    return {f'{prefix}{k}': v for k, v in d.items()}


def flatten_dict(d: typing.Mapping[str, Any]) -> dict[str, Any]:
    """将嵌套的字典展平为单层字典

    Example
    -------
    ```py
    dic = {
        'loss': 1,
        'avg': { 'loss': 1.2 },
        'Exp': { 'Round0': { 'loss': 2 }, 'Round1': { 'loss': 1 } }
    }
    flatten_dict(dic)
    # {'loss': 1, 'avg.loss': 1.2, 'Exp.Round0.loss': 2, 'Exp.Round1.loss': 1}
    ```
    """
    result: dict[str, Any] = {}

    def recurse(sub_dict: typing.Mapping[str, Any], parent_key: str = ''):
        for k, v in sub_dict.items():
            new_key = f'{parent_key}.{k}' if parent_key else k

            if isinstance(v, dict):
                recurse(v, new_key)
            else:
                result[new_key] = v

    recurse(d)
    return result


# ==================== Augmented Dict ====================


class DottableDict(dict):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._convert_nested_dicts()

    def _convert_nested_dicts(self):
        for key, value in list(self.items()):
            if isinstance(value, dict) and not isinstance(value, DottableDict):
                super().__setitem__(key, DottableDict(value))

    def __getattr__(self, key):
        try:
            return self[key]
        except KeyError:
            raise AttributeError(f"'DottableDict' object has no attribute '{key}'")

    def __setattr__(self, key, value):
        if key.startswith('_'):
            super().__setattr__(key, value)
        else:
            self[key] = value

    def _process_nested_key(self, key):
        """处理可能包含点号的键

        Returns:
            tuple: (is_nested, keys)
                - is_nested: 是否是嵌套的键
                - keys: 键的列表
        """
        if not isinstance(key, str):
            return False, None
        if '.' not in key:
            return False, None
        return True, key.split('.')

    def __getitem__(self, key):
        is_nested, keys = self._process_nested_key(key)
        if not is_nested:
            return super().__getitem__(key)
        if not keys:
            return super().__getitem__(key)

        d = self
        for k in keys:
            d = d[k]
        return d

    def __setitem__(self, key, value):
        if isinstance(value, dict) and not isinstance(value, DottableDict):
            value = DottableDict(value)

        is_nested, keys = self._process_nested_key(key)
        if not is_nested or not keys:
            super().__setitem__(key, value)
            return

        d = self
        for k in keys[:-1]:
            if k not in d:
                d[k] = DottableDict()
            d = d[k]
        d[keys[-1]] = value

    def __delattr__(self, key):
        try:
            del self[key]
        except KeyError:
            raise AttributeError(f"'DottableDict' object has no attribute '{key}'")

    def __delitem__(self, key):
        is_nested, keys = self._process_nested_key(key)
        if not is_nested or not keys:
            super().__delitem__(key)
            return

        d = self
        for k in keys[:-1]:
            if k not in d:
                raise KeyError(key)
            d = d[k]
        del d[keys[-1]]

    def __contains__(self, key):
        is_nested, keys = self._process_nested_key(key)
        if not is_nested or not keys:
            return super().__contains__(key)

        try:
            d = self
            for k in keys:
                d = d[k]
            return True
        except (KeyError, TypeError):
            return False

    def update(self, other: typing.Mapping | None = None, **kwargs):
        """
        Recursively update the dictionary with items from another dictionary object and/or from an iterable of key/value pairs.
        """
        items_to_update = {}
        if other:
            items_to_update.update(other)
        if kwargs:
            items_to_update.update(kwargs)

        for k, v in items_to_update.items():
            if isinstance(v, dict):
                # If key exists and is a DottableDict, update it recursively.
                # Otherwise, create a new DottableDict.
                if k in self and isinstance(self.get(k), DottableDict):
                    self[k].update(v)
                else:
                    self[k] = DottableDict(v)
            else:
                self[k] = v


# ==================== 同构字典列表操作 ====================
# 本质上其实就是实现以 list[dict] 为组织形式的异形数据库的操作; 每个 dict 都是一个数据项


def aggregate_dicts(
    fn_aggerate: Callable[[str, typing.Sequence[Any]], Any],
    dicts: typing.Sequence[typing.Mapping],
) -> dict:
    """聚合若干个同构的字典; 可以是复杂嵌套的结构

    输入多个同构的字典, 通过一个聚合函数合并成一个字典, 例如

    ```py
    def fn_aggerate(key: str, values: list[Any]) -> Any:
        return sum(values) / len(values)
    ```

    递归遍历:

    - 当 value 为 dict 的时候, 递归调用 `aggregate_dicts`
    - 其他情况（包括 list np.array 等可迭代类型） 都视为原子结构，被放入 List 中进行聚合

    Parameters
    ----------
    - `fn_aggerate` : `Callable[[str, list[Any]], Any]`
        - 聚合函数，接收键名和值列表，返回聚合后的值
    - `dicts` : `list[dict]`
        - 要聚合的字典列表，应该具有相同的键结构

    Returns
    ----------
    - `dict`
        - 聚合后的字典
    """
    if not dicts:
        return {}

    if not isinstance(dicts, list):
        raise ValueError(f'传入的 dicts 必须是一个列表, 而非 {type(dicts).__name__}')

    if len(dicts) == 1:
        return typing.cast(Any, dicts[0]).copy()

    # 类型检查：确保所有元素都是字典
    for i, d in enumerate(dicts):
        if not isinstance(d, dict):
            raise TypeError(
                f'Element at index {i} is not a dict, got {type(d).__name__}: {repr(d)}'
            )

    # 获取所有字典的键的并集
    all_keys = set()
    for d in dicts:
        all_keys.update(d.keys())

    result = {}

    for key in all_keys:
        # 收集该键在所有字典中的值
        values = []
        for d in dicts:
            if key in d:
                values.append(d[key])

        if not values:
            continue

        # 检查是否所有值都是字典类型
        if all(isinstance(v, dict) for v in values):
            # 递归聚合嵌套字典
            result[key] = aggregate_dicts(fn_aggerate, values)
        else:
            # 使用聚合函数处理非字典值
            result[key] = fn_aggerate(key, values)

    return result


def diff_dicts(
    dicts: typing.Sequence[typing.Mapping],
    heterogeneous_key: Literal['skip', 'diff', 'raise'] = 'skip',
):
    """输入同构的字典，返回他们存在差异性的部分

    比较多个字典，返回其中存在差异的键值对。对于每个键，如果所有字典的值都相同，
    则该键不会包含在结果中；如果值存在差异，则保留在结果中。

    Parameters
    ----------
    - `dicts` : `list[dict]`
        - 要比较的字典列表，应该具有相同或相似的键结构
    - `heterogeneous_key` : `Literal['skip', 'diff', 'raise']`, optional
        - 处理异构键（只在部分字典中存在的键）的策略, by default 'skip'
        - 'skip': 跳过异构键，不包含在结果中
        - 'diff': 将异构键视为差异，包含在结果中
        - 'raise': 遇到异构键时抛出异常

    Returns
    ----------
    - `dict`
        - 包含存在差异的键值对的字典，每个值是一个列表，包含所有输入字典中该键的值
    """
    if not dicts:
        return {}

    if not isinstance(dicts, list):
        raise ValueError(f'传入的 dicts 必须是一个列表, 而非 {type(dicts).__name__}')

    if len(dicts) <= 1:
        return {}

    # 类型检查：确保所有元素都是字典
    for i, d in enumerate(dicts):
        if not isinstance(d, dict):
            raise TypeError(
                f'Element at index {i} is not a dict, got {type(d).__name__}: {repr(d)}'
            )

    # 获取所有字典的键的并集和交集
    all_keys = set()
    for d in dicts:
        all_keys.update(d.keys())

    # 检查每个键在所有字典中的存在情况
    common_keys = set(dicts[0].keys())
    for d in dicts[1:]:
        common_keys.intersection_update(d.keys())

    heterogeneous_keys = all_keys - common_keys

    # 处理异构键
    if heterogeneous_keys:
        if heterogeneous_key == 'raise':
            raise ValueError(
                f'发现异构键: {heterogeneous_keys}。这些键不在所有字典中都存在。'
            )
        elif heterogeneous_key == 'skip':
            # 只处理公共键
            keys_to_process = common_keys
        elif heterogeneous_key == 'diff':
            # 处理所有键，异构键自动视为差异
            keys_to_process = all_keys
        else:
            raise ValueError(f'不支持的 heterogeneous_key 值: {heterogeneous_key}')
    else:
        keys_to_process = all_keys

    result = {}

    for key in keys_to_process:
        # 收集该键在所有字典中的值
        values = []
        for d in dicts:
            if key in d:
                values.append(d[key])
            else:
                # 如果键不存在，添加 None 作为占位符（仅在 diff 模式下可能发生）
                values.append(None)

        # 检查是否所有值都是字典类型
        if all(isinstance(v, dict) or v is None for v in values):
            # 过滤掉 None 值进行递归比较
            non_none_values = [v for v in values if v is not None]
            if len(set(id(v) for v in non_none_values)) > 1:  # 有不同的字典对象
                # 递归比较嵌套字典
                nested_diff = diff_dicts(non_none_values, heterogeneous_key)
                if nested_diff:  # 只有在存在差异时才添加
                    result[key] = nested_diff
            # 如果有 None 值且处于 diff 模式，也视为差异
            elif None in values and heterogeneous_key == 'diff':
                result[key] = values
        else:
            # 比较非字典值
            # 检查是否所有值都相同
            unique_values = set()
            for v in values:
                if isinstance(v, (list, dict)):
                    # 对于可变类型，使用字符串表示进行比较
                    unique_values.add(str(v))
                else:
                    unique_values.add(v)

            if len(unique_values) > 1:
                # 存在差异，保留在结果中
                result[key] = values

    return result


def groupby_dicts(
    groupby: str | Callable[[dict], str | int],
    dicts: typing.Sequence[dict],
    fn_process: Callable[[str, typing.Sequence[dict]], Any] | None = None,
) -> dict:
    """类似对数据表格的 groupby 操作; 只不过这里使用 list[dict] 来表达复杂数据集合

    将字典列表按照指定的键或函数进行分组，返回分组后的结果字典。

    Parameters
    ----------
    - `groupby` : `str | Callable[[dict], str | int]`
        - 分组依据，可以是字典的键名（字符串）或者是一个函数，函数接收字典并返回分组标识
    - `dicts` : `list[dict]`
        - 要进行分组操作的字典列表
    - `fn_process` : `Callable[[str, list[dict]], Any]`, optional
        - 对每个分组进行处理的函数，接收分组标识和该组的字典列表，返回处理结果, by default None

    Returns
    ----------
    - `dict`
        - 分组后的结果字典，键为分组标识，值为该组的字典列表（或处理后的结果）
    """
    if not dicts:
        return {}

    if not isinstance(dicts, list):
        raise ValueError(f'传入的 dicts 必须是一个列表, 而非 {type(dicts).__name__}')

    # 类型检查：确保所有元素都是字典
    for i, d in enumerate(dicts):
        if not isinstance(d, dict):
            raise TypeError(
                f'Element at index {i} is not a dict, got {type(d).__name__}: {repr(d)}'
            )

    # 确定分组函数
    if isinstance(groupby, str):
        # 如果是字符串，创建一个获取该键值的函数
        key_func = lambda d: d.get(groupby)
    elif callable(groupby):
        key_func = groupby
    else:
        raise TypeError(
            f'groupby 参数必须是字符串或可调用对象，而非 {type(groupby).__name__}'
        )

    # 进行分组
    groups = {}
    for d in dicts:
        try:
            group_key = key_func(d)
            # 将 None 转换为字符串以便作为字典键
            if group_key is None:
                group_key = 'None'
            # 确保键是可哈希的
            elif not isinstance(group_key, (str, int, float, tuple)):
                group_key = str(group_key)

            if group_key not in groups:
                groups[group_key] = []
            groups[group_key].append(d)
        except Exception as e:
            raise ValueError(f'分组函数执行失败，字典: {d}, 错误: {e}')

    # 如果提供了处理函数，对每个分组进行处理
    if fn_process is not None:
        if not callable(fn_process):
            raise TypeError('fn_process 必须是可调用对象')

        result = {}
        for group_key, group_dicts in groups.items():
            try:
                ans = fn_process(group_key, group_dicts)
                if ans is not None:
                    result[group_key] = ans
            except Exception as e:
                raise ValueError(f'处理函数执行失败，分组: {group_key}, 错误: {e}')
        return result

    return groups
