import typing
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import List, Optional, Set, TypedDict, Union
from xml.etree import ElementTree as ET

from zyplib.annotation.event import Event


def _require_xml_text(element: ET.Element, tag: str) -> str:
    child = element.find(tag)
    if child is None or child.text is None:
        raise ValueError(f'Missing XML field: {tag}')
    return child.text


@dataclass
class ScoredEventType:
    type: int
    name: str

    def __hash__(self) -> int:
        return self.type

    def __eq__(self, __value: object) -> bool:
        if isinstance(__value, ScoredEventType):
            return self.type == __value.type
        elif isinstance(__value, int):
            return self.type == __value
        elif isinstance(__value, str):
            return self.name == __value
        else:
            return False


@dataclass
class ScoredEvent(ScoredEventType):
    type: int
    name: str
    time: float
    duration: float
    manully_scored: int
    param1: Optional[int] = None
    param2: Optional[int] = None
    param3: Optional[int] = None

    def same_type(self, __value: Union[int, str, ScoredEventType]) -> bool:
        return super().__eq__(__value)

    def to_event(self) -> Event:
        return Event(name=self.name, onset=self.time, duration=self.duration)


class PSGScoredEvents:
    """CMPPSGSCOREDATA XML 文件的读取"""

    def __init__(self) -> None:
        self.base_name: str | None = None  # 如 'cyq202107220206' 包含了姓名和基本信息
        self.created_on: datetime | None = None
        self.scored_events: List[ScoredEvent] = []
        self.sleep_stages: List[int] = []

    def push_event_(self, onset, duration, name, type=-1, manually_scored=-1):
        """添加事件"""
        scored_event = ScoredEvent(
            type=type,
            name=name,
            time=onset,
            duration=duration,
            manully_scored=manually_scored,
        )
        self.scored_events.append(scored_event)

    @property
    def unique_events(self) -> Set[ScoredEventType]:
        """获取所有的唯一的事件"""
        return set(
            [
                ScoredEventType(type=score_event.type, name=score_event.name)
                for score_event in self.scored_events
            ]
        )

    def filter(self, *event_types: Union[int, str, ScoredEventType]) -> 'PSGScoredEvents':
        """根据事件类型过滤 scored_events, 并创建一个新的 PSGScoredData

        根据输入的 event_type 来过滤 scored_events, 并创建一个新的 PSGScoredData

        Parameters
        ----------
        - `event_types`: 事件类型, 可以是 int, str, ScoredEventType

        Returns
        ---------
        - `PSGScoredData`: 过滤后的 PSGScoredData
        """
        selected_event_types = set(event_types)
        new_scored_events = []
        for score_event in self.scored_events:
            for event_type in selected_event_types:
                if score_event.same_type(event_type):
                    # print('Same type', score_event)
                    new_scored_events.append(score_event)
        psg_scored_data = PSGScoredEvents()
        psg_scored_data.created_on = self.created_on
        psg_scored_data.scored_events = new_scored_events
        psg_scored_data.sleep_stages = self.sleep_stages
        return psg_scored_data

    def filter_(self, *event_types: Union[int, str, ScoredEventType]):
        """据事件类型过滤 scored_events, 同 `filter`, 但是直接更改 self"""
        psg_scored_data = self.filter(*event_types)
        self.created_on = psg_scored_data.created_on
        self.scored_events = psg_scored_data.scored_events
        self.sleep_stages = psg_scored_data.sleep_stages

    @staticmethod
    def from_xml_file(xml_fpath: str):
        fname = Path(xml_fpath).name
        base_name = fname.split('.')[0]

        # 1. 读取 xml 文件
        tree = ET.parse(xml_fpath)
        # 2. 读取 /SCOREDEVENTS/SCOREDEVENT
        root = tree.getroot()

        score_events = []
        SCOREDEVENTS = typing.cast(ET.Element, root.find('SCOREDEVENTS'))
        for scoredevent in SCOREDEVENTS.iter('SCOREDEVENT'):
            score_event = ScoredEvent(
                type=int(_require_xml_text(scoredevent, 'TYPE')),
                name=_require_xml_text(scoredevent, 'NAME'),
                time=float(_require_xml_text(scoredevent, 'TIME')),
                duration=float(_require_xml_text(scoredevent, 'DURATION')),
                manully_scored=int(_require_xml_text(scoredevent, 'MANUALLYSCORED')),
                param1=int(_require_xml_text(scoredevent, 'PARAM1')),
                param2=int(_require_xml_text(scoredevent, 'PARAM2')),
                param3=int(_require_xml_text(scoredevent, 'PARAM3')),
            )
            score_events.append(score_event)
        # 3. 读取 /SLEEPSTAGES/SLEEPSTAGE
        sleep_stages = []
        for sleepstage in root.findall('SLEEPSTAGES/SLEEPSTAGE'):
            sleep_stages.append(int(typing.cast(str, sleepstage.text)))
        # 4. 返回
        psg_scored_data = PSGScoredEvents()
        created_on: datetime = datetime.strptime(
            _require_xml_text(root, 'CREATEDON'), '%Y/%m/%d %H:%M:%S'
        )
        psg_scored_data.created_on = created_on
        psg_scored_data.scored_events = score_events
        psg_scored_data.sleep_stages = sleep_stages
        psg_scored_data.base_name = base_name
        return psg_scored_data


def read_scoreddata_xml(xml_fpath: str) -> List[Event]:
    psg_events = PSGScoredEvents.from_xml_file(xml_fpath)
    return [event.to_event() for event in psg_events.scored_events]
