from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import Any, TYPE_CHECKING

if TYPE_CHECKING:
    from yuxi.channels.models import ChannelAccountSnapshot


@dataclass
class ChatAbortEntry:
    task: Any = None
    event: asyncio.Event = field(default_factory=asyncio.Event)

    def abort(self) -> None:
        self.event.set()
        if self.task is not None and not self.task.done():
            self.task.cancel()

    def is_aborted(self) -> bool:
        return self.event.is_set()


@dataclass
class ChatRunBuffer:
    run_id: str = ""
    _chunks: list[str] = field(default_factory=list)
    _finished: bool = False

    def append_chunk(self, chunk: str) -> None:
        self._chunks.append(chunk)

    def get_full_text(self) -> str:
        return "".join(self._chunks)

    def get_chunks(self) -> list[str]:
        return list(self._chunks)

    def is_finished(self) -> bool:
        return self._finished

    def mark_finished(self) -> None:
        self._finished = True

    def chunk_count(self) -> int:
        return len(self._chunks)


@dataclass
class GatewayRequestContext:
    runtime_config: dict = field(default_factory=dict)
    runtime_snapshot: ChannelAccountSnapshot | None = None

    start_channel: Callable[[str, dict], Awaitable[None]] | None = None
    stop_channel: Callable[[str], Awaitable[None]] | None = None
    mark_channel_logged_out: Callable[[str, str], Awaitable[None]] | None = None

    broadcast_fn: Callable[[str, Any], Awaitable[None]] | None = None
    node_send_to_session_fn: Callable[[str, str, Any], Awaitable[None]] | None = None

    get_snapshot: Callable[[str], Awaitable[ChannelAccountSnapshot]] | None = None

    dedupe: Any = None

    chat_abort_controllers: dict[str, ChatAbortEntry] = field(default_factory=dict)
    chat_run_buffers: dict[str, ChatRunBuffer] = field(default_factory=dict)

    def get_runtime_config(self, key: str, default=None):
        return self.runtime_config.get(key, default)

    def set_runtime_config(self, key: str, value) -> None:
        self.runtime_config[key] = value

    def add_chat_run(self, run_id: str, entry: ChatAbortEntry, buffer: ChatRunBuffer | None = None) -> None:
        self.chat_abort_controllers[run_id] = entry
        if buffer is not None:
            self.chat_run_buffers[run_id] = buffer

    def remove_chat_run(self, run_id: str) -> None:
        self.chat_abort_controllers.pop(run_id, None)
        self.chat_run_buffers.pop(run_id, None)

    def abort_chat_run(self, run_id: str) -> bool:
        entry = self.chat_abort_controllers.get(run_id)
        if entry is None:
            return False
        entry.abort()
        return True

    def get_chat_buffer(self, run_id: str) -> ChatRunBuffer | None:
        return self.chat_run_buffers.get(run_id)

    async def broadcast(self, event: str, payload: Any = None) -> None:
        if self.broadcast_fn is not None:
            await self.broadcast_fn(event, payload)

    async def node_send_to_session(self, session_id: str, event: str, payload: Any = None) -> None:
        if self.node_send_to_session_fn is not None:
            await self.node_send_to_session_fn(session_id, event, payload)