from __future__ import annotations

import abc
from contextlib import contextmanager
from dataclasses import dataclass
from enum import Enum
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union

if TYPE_CHECKING:
    from taskweaver.memory.attachment import Attachment, AttachmentType
    from taskweaver.memory.post import Post
    from taskweaver.memory.type_vars import RoleName


class EventScope(Enum):
    session = "session"
    round = "round"
    post = "post"


class SessionEventType(Enum):
    session_start = "session_start"
    session_end = "session_end"
    session_new_round = "session_new_round"


class RoundEventType(Enum):
    round_start = "round_start"
    round_end = "round_end"
    round_error = "round_error"
    round_new_post = "round_new_post"


class PostEventType(Enum):
    post_start = "post_start"
    post_end = "post_end"
    post_error = "post_error"

    post_status_update = "post_status_update"
    post_send_to_update = "post_send_to_update"
    post_message_update = "post_message_update"
    post_attachment_update = "post_attachment_update"


@dataclass
class TaskWeaverEvent:
    scope: EventScope
    t: Union[SessionEventType, RoundEventType, PostEventType]
    round_id: Optional[str]
    post_id: Optional[str]
    msg: str
    extra: Any = None


class SessionEventHandler(abc.ABC):
    @abc.abstractmethod
    def handle(self, event: TaskWeaverEvent):
        pass


class SessionEventHandlerBase(SessionEventHandler):
    def handle(self, event: TaskWeaverEvent):
        if event.scope == EventScope.session:
            assert isinstance(event.t, SessionEventType)
            session_event_type: SessionEventType = event.t
            self.handle_session(
                session_event_type,
                event.msg,
                event.extra,
            )
        elif event.scope == EventScope.round:
            assert isinstance(event.t, RoundEventType)
            assert event.round_id is not None
            round_event_type: RoundEventType = event.t
            self.handle_round(
                round_event_type,
                event.msg,
                event.extra,
                event.round_id,
            )

        elif event.scope == EventScope.post:
            assert isinstance(event.t, PostEventType)
            assert event.post_id is not None
            assert event.round_id is not None
            post_event_type: PostEventType = event.t
            self.handle_post(
                post_event_type,
                event.msg,
                event.extra,
                event.post_id,
                event.round_id,
            )

    def handle_session(
        self,
        type: SessionEventType,
        msg: str,
        extra: Any,
        **kwargs: Any,
    ):
        pass

    def handle_round(
        self,
        type: RoundEventType,
        msg: str,
        extra: Any,
        round_id: str,
        **kwargs: Any,
    ):
        pass

    def handle_post(
        self,
        type: PostEventType,
        msg: str,
        extra: Any,
        post_id: str,
        round_id: str,
        **kwargs: Any,
    ):
        pass


class PostEventProxy:
    def __init__(self, emitter: SessionEventEmitter, round_id: str, post: Post) -> None:
        self.emitter = emitter
        self.round_id = round_id
        self.post = post
        self.message_is_end = False
        self.create("Post created")

    def create(self, message: str):
        self._emit(
            PostEventType.post_start,
            message,
            {
                "role": self.post.send_from,
            },
        )

    def update_send_to(self, send_to: RoleName):
        self.post.send_to = send_to
        self._emit(
            PostEventType.post_send_to_update,
            "",
            {
                "role": send_to,
            },
        )

    def update_status(self, status: str):
        self._emit(PostEventType.post_status_update, status)

    def update_message(self, message: str, is_end: bool = True):
        assert not self.message_is_end, "Cannot update message when update is finished"
        self.post.message += message
        self.message_is_end = is_end
        self._emit(
            PostEventType.post_message_update,
            message,
            {"is_end": is_end},
        )

    def update_attachment(
        self,
        message: str,
        type: Optional[AttachmentType] = None,
        extra: Any = None,
        id: Optional[str] = None,
        is_end: bool = True,
    ) -> Attachment:
        from taskweaver.memory.attachment import Attachment

        if id is not None:
            attachment = self.post.attachment_list[-1]
            assert id == attachment.id
            if type is not None:
                assert type == attachment.type
            attachment.content += message
            attachment.extra = extra
        else:
            assert type is not None, "type is required when creating new attachment"
            attachment = Attachment.create(
                type=type,
                content=message,
                extra=extra,
                id=id,
            )
            self.post.add_attachment(attachment)
        self._emit(
            PostEventType.post_attachment_update,
            message,
            {
                "type": type,
                "extra": extra,
                "id": attachment.id,
                "is_end": is_end,
            },
        )
        return attachment

    def error(self, msg: str):
        self.post.attachment_list = []
        self.post.message = msg
        self._emit(PostEventType.post_error, msg)

    def end(self, msg: str = ""):
        self._emit(PostEventType.post_end, msg)
        return self.post

    def _emit(
        self,
        event_type: PostEventType,
        message: str,
        extra: Dict[str, Any] = {},
    ):
        self.emitter.emit(
            TaskWeaverEvent(
                EventScope.post,
                event_type,
                self.round_id,
                self.post.id,
                message,
                extra=extra,
            ),
        )


class SessionEventEmitter:
    def __init__(self):
        self.handlers: List[SessionEventHandler] = []
        self.current_round_id: Optional[str] = None

    def emit(self, event: TaskWeaverEvent):
        for handler in self.handlers:
            handler.handle(event)

    def start_round(self, round_id: str):
        self.current_round_id = round_id
        self.emit(
            TaskWeaverEvent(
                EventScope.round,
                RoundEventType.round_start,
                round_id,
                None,
                "Round started",
            ),
        )

    def create_post_proxy(self, send_from: RoleName) -> PostEventProxy:
        assert self.current_round_id is not None, "Cannot create post proxy without a round in active"
        from taskweaver.memory.post import Post

        return PostEventProxy(
            self,
            self.current_round_id,
            Post.create(message="", send_from=send_from),
        )

    def emit_error(self, msg: str):
        self.emit(
            TaskWeaverEvent(
                EventScope.round,
                RoundEventType.round_error,
                self.current_round_id,
                None,
                msg,
            ),
        )

    def end_round(self, round_id: str):
        assert self.current_round_id == round_id, "Cannot end round that is not in active"
        self.emit(
            TaskWeaverEvent(
                EventScope.round,
                RoundEventType.round_end,
                round_id,
                None,
                "Round ended",
            ),
        )
        self.current_round_id = None

    def register(self, handler: SessionEventHandler):
        self.handlers.append(handler)

    def unregister(self, handler: SessionEventHandler):
        self.handlers.remove(handler)

    @contextmanager
    def handle_events_ctx(self, handler: Optional[SessionEventHandler] = None):
        if handler is None:
            yield
        else:
            self.register(handler)
            yield
            self.unregister(handler)
