"""
自定义一种基于 HDF5 的脑电信号数据存储格式，作为标准存储格式
"""

import re
import typing
from typing import TypedDict
from pathlib import Path
from typing import Union, overload

import h5py
import numpy as np

from zyplib.annotation.event import Event, events_to_dict

from ._base import BaseEEGSample, EEGSampleInterface
from .utils import match_channels_regex


class NormalizedEEGData(TypedDict):
    data_list: list[np.ndarray]
    channels: list[str]
    srate_list: list[float]
    time_length: float
    events: list[Event]
    description: str


def _to_string_list(values: typing.Any) -> list[str]:
    result: list[str] = []
    for value in typing.cast(typing.Iterable[typing.Any], values):
        if isinstance(value, bytes):
            result.append(value.decode('utf-8'))
        else:
            result.append(str(value))
    return result


def _read_h5_event_names(dataset: h5py.Dataset) -> list[str]:
    values = dataset[()]
    if isinstance(values, np.ndarray):
        return _to_string_list(values.tolist())
    return _to_string_list(values)


def _normalize_eeg_data_for_h5(
    data: Union[np.ndarray, typing.Sequence[np.ndarray]],
    channels: Union[typing.Sequence[str], np.ndarray],
    srate: Union[float, int, typing.Sequence[float], typing.Sequence[int], np.ndarray],
    events: typing.Sequence[Event] | None = None,
    description: str = '',
) -> NormalizedEEGData:
    """将脑电数据标准化为 H5 存储格式

    完全按照原始 _as_std_type_ 的逻辑进行标准化处理。

    Parameters
    ----------
    - `data` : `Union[np.ndarray, List[np.ndarray]]`
        - 脑电信号数据，可以是 2D ndarray 或 list of 1D ndarray
    - `channels` : `Union[List[str], np.ndarray]`
        - 通道名称，可以是 list 或 ndarray
    - `srate` : `Union[float, int, List[float], List[int], np.ndarray]`
        - 采样率，可以是单个值、列表或数组
    - `events` : `list[Event]`, optional
        - 事件标注列表
    - `description` : `str`, optional
        - 描述信息

    Returns
    ----------
    - `dict`
        - 标准化后的数据字典
    """
    events_list = list(events or [])

    # 1. 处理 Channels：确保是 list
    if isinstance(channels, np.ndarray):
        channels = _to_string_list(channels.tolist())
    else:
        channels = [str(channel) for channel in channels]

    # 2. 处理 Srate：确保是 list of float
    if isinstance(srate, np.ndarray):
        srate = typing.cast(list[float], srate.squeeze().astype(float).tolist())
    if isinstance(srate, (int, float)):
        srate = [float(srate)] * len(channels)
    else:
        srate = [float(value) for value in srate]

    # 3. 处理 Data：确保是 list of float32 arrays
    if isinstance(data, np.ndarray):
        # 如果是 2D ndarray，转换为 list of 1D arrays
        if data.ndim == 2:
            # 验证通道数匹配
            if data.shape[0] != len(channels):
                raise ValueError(
                    f'Data has {data.shape[0]} channels but {len(channels)} channel names provided'
                )
            data_list = [data[i, :].astype(np.float32) for i in range(data.shape[0])]
        else:
            raise ValueError(f'Unsupported ndarray dimension: {data.ndim}')
    elif (
        isinstance(data, typing.Sequence)
        and len(data) > 0
        and isinstance(data[0], np.ndarray)
    ):
        # 如果已经是 list of ndarray，确保每个都是 float32
        # 验证通道数匹配
        if len(data) != len(channels):
            raise ValueError(
                f'Data has {len(data)} channels but {len(channels)} channel names provided'
            )
        data_list = []
        for i in range(len(channels)):
            channel_data = data[i].astype(np.float32)
            data_list.append(channel_data)
    else:
        raise ValueError(f'Unsupported data type: {type(data)}')

    # 4. 计算 TimeLength
    if data_list and len(data_list) > 0:
        time_length = len(data_list[0]) / srate[0]
    else:
        time_length = -1

    return {
        'data_list': data_list,
        'channels': channels,
        'srate_list': srate,
        'time_length': time_length,
        'events': events_list,
        'description': str(description),
    }


@overload
def save_eeg_hdf5(
    path: str | Path,
    eeg_sample: EEGSampleInterface,
    data: None = None,
    channels: None = None,
    srate: None = None,
    events: typing.Sequence[Event] | None = None,
    description: str = '',
    **kwargs: typing.Any,
) -> str: ...


@overload
def save_eeg_hdf5(
    path: str | Path,
    eeg_sample: None = None,
    data: np.ndarray | None = None,
    channels: typing.Sequence[str] | None = None,
    srate: float | None = None,
    events: typing.Sequence[Event] | None = None,
    description: str = '',
    **kwargs: typing.Any,
) -> str: ...


def save_eeg_hdf5(
    path: str | Path,
    eeg_sample: EEGSampleInterface | None = None,
    data: np.ndarray | None = None,
    channels: typing.Sequence[str] | None = None,
    srate: float | None = None,
    events: typing.Sequence[Event] | None = None,
    description: str = '',
    **kwargs,
) -> str:
    """保存脑电信号为标准 H5 格式文件

    支持两种调用方式：
    1. 传入 EEGSampleInterface 对象
    2. 传入独立的 data, channels, srate 等参数

    Parameters
    ----------
    - `path` : `str`
        - 保存文件的路径
    - `eeg_sample` : `EEGSampleInterface`, optional
        - 脑电信号样本对象
    - `data` : `np.ndarray`, optional
        - 脑电信号数据，形状为 (n_channels, n_samples)
    - `channels` : `list[str]`, optional
        - 通道名称列表
    - `srate` : `float`, optional
        - 采样率
    - `events` : `list[Event]`, optional
        - 事件标注列表
    - `description` : `str`, optional
        - 文件描述信息

    Examples
    ----------
    ```py
    >>> # 方式1：使用 EEGSampleInterface 对象
    >>> save_eegh5_file('output.h5', eeg_sample=my_eeg_sample)

    >>> # 方式2：使用独立参数
    >>> save_eegh5_file('output.h5',
    ...                 data=data_array,
    ...                 channels=['C3', 'C4'],
    ...                 srate=250.0,
    ...                 events=event_list)
    ```
    """
    # 参数解析
    if eeg_sample is not None:
        # 使用 EEGSampleInterface 对象
        data = eeg_sample.data
        channels = eeg_sample.channels
        srate = eeg_sample.srate
        events = events or eeg_sample.events
        description = description or getattr(eeg_sample, 'description', '')
    else:
        # 使用独立参数
        if data is None or channels is None or srate is None:
            raise ValueError(
                'When eeg_sample is not provided, data, channels, and srate are required'
            )
        events = events or []

    # 标准化数据
    normalized_data = _normalize_eeg_data_for_h5(
        data, channels, srate, events, description
    )

    # 确保路径以 .eeg.h5 结尾
    path = str(path)
    if not path.endswith('.h5'):
        path += '.eeg.h5'
    elif not path.endswith('.eeg.h5'):
        path = path[:-3] + '.eeg.h5'

    # 保存到 H5 文件
    with h5py.File(path, 'w') as f:
        # 保存元数据
        meta_group = f.create_group('meta')
        meta_group.attrs['Channels'] = normalized_data['channels']
        meta_group.attrs['Srate'] = normalized_data['srate_list']
        meta_group.attrs['TimeLength'] = normalized_data['time_length']
        meta_group.attrs['Description'] = normalized_data['description']

        # 保存数据
        data_group = f.create_group('Data')
        for i, channel_data in enumerate(normalized_data['data_list']):
            data_group.create_dataset(str(i), data=channel_data)

        # 保存事件
        events_group = f.create_group('Events')
        if normalized_data['events']:
            events_dict = events_to_dict(normalized_data['events'])
            events_group.create_dataset('Onset', data=events_dict['Onset'])
            events_group.create_dataset('Duration', data=events_dict['Duration'])
            # 将字符串转换为字节串以兼容 H5py
            event_names = [name.encode('utf-8') for name in events_dict['Name']]
            events_group.create_dataset('Name', data=event_names)
    return path


def load_eeg_hdf5(
    path: str | Path,
    meta_only: bool = False,
    match_channels: typing.Sequence[Union[str, re.Pattern]] | None = None,
) -> BaseEEGSample:
    """加载标准 H5 格式脑电文件

    Parameters
    ----------
    - `path` : `str`
        - H5 文件的路径
    - `meta_only` : `bool`, optional
        - 是否只读取元数据, by default False
    - `match_channels` : `List[Union[str, re.Pattern]]`, optional
        - 正则表达式列表，用于匹配特定通道, by default None

    Returns
    ----------
    - `BaseEEGSample`
        - 脑电信号样本对象

    Examples
    ----------
    ```py
    >>> # 加载完整数据
    >>> eeg_sample = load_eegh5_file('data.eeg.h5')

    >>> # 只加载元数据
    >>> eeg_meta = load_eegh5_file('data.eeg.h5', meta_only=True)

    >>> # 匹配特定通道
    >>> eeg_filtered = load_eegh5_file('data.eeg.h5', match_channels=[r'C[34]'])
    ```
    """
    # 验证文件存在
    if not Path(path).exists():
        raise FileNotFoundError(f'File not found: {path}')

    with h5py.File(path, 'r') as f:
        # 读取元数据
        meta = typing.cast(h5py.Group, f['meta'])
        channels = _to_string_list(meta.attrs['Channels'].tolist())  # pyright: ignore[reportAttributeAccessIssue]
        srate_list = [float(value) for value in meta.attrs['Srate'].tolist()]  # pyright: ignore[reportAttributeAccessIssue]
        time_length = float(meta.attrs['TimeLength'])  # pyright: ignore[reportArgumentType]
        description = str(meta.attrs['Description'])

        # 读取事件
        events: list[Event] = []
        events_group = typing.cast(h5py.Group, f['Events'])
        if len(events_group.keys()) > 0:
            onset_values = np.atleast_1d(events_group['Onset'][()])  # type: ignore
            duration_values = np.atleast_1d(events_group['Duration'][()])  # type: ignore
            name_values = _read_h5_event_names(
                typing.cast(h5py.Dataset, events_group['Name'])
            )
            # 处理多个事件
            if len(onset_values) > 0:
                for i, event_name in enumerate(name_values):
                    events.append(
                        Event(
                            onset=float(onset_values[i]),
                            duration=float(duration_values[i]),
                            name=event_name,
                        )
                    )

        # 处理通道匹配
        if match_channels is not None:
            matched_ch_names, matched_ch_index = match_channels_regex(
                channels, match_channels
            )
            channels = matched_ch_names
            srate_list = [srate_list[i] for i in matched_ch_index]
            include_idx = matched_ch_index
        else:
            include_idx = list(range(len(channels)))

        # 读取数据
        if meta_only:
            # 创建空数据数组
            data = np.empty((len(channels), 0), dtype=np.float32)
        else:
            data_group = typing.cast(h5py.Group, f['Data'])
            data_list: list[np.ndarray] = []
            for i in include_idx:
                data_list.append(np.asarray(data_group[str(i)][()], dtype=np.float32))  # pyright: ignore[reportIndexIssue]
            data = np.array(data_list, dtype=np.float32)

    # 获取采样率（假设所有通道采样率相同）
    # srate = srate_list[0] if srate_list else 1.0
    srate = srate_list[0] if np.allclose(srate_list, srate_list[0]) else srate_list

    # 创建 BaseEEGSample 对象
    return BaseEEGSample(
        data=data,
        srate=typing.cast(float, srate),
        channels=channels,
        name=Path(path).stem,
        filepath=Path(path),
        units='uV',
        meta={
            'time_length': time_length,
            'description': description,
            'srate_list': srate_list,
        },
        events=events,
        duration=time_length,
    )
