from __future__ import annotations

import asyncio
import logging
import time
from collections.abc import Awaitable, Callable

from yuxi.channel.runtime.backoff import BackoffConfig, ErrorBackoff
from yuxi.channel.config.defaults import TIMEOUT
from yuxi.channel.streaming.directives import (
    DirectiveHandler,
    StreamingDirective,
    StreamingDirectiveParser,
)
from yuxi.channel.streaming.models import DraftStreamSession, DraftStreamState

logger = logging.getLogger(__name__)

SendOrEditFn = Callable[[str, str | None], Awaitable[str | None]]
ProgressCallback = Callable[[str], Awaitable[None]]

_MAX_FLUSH_ITERATIONS = 100
_TOOL_STATUS_PREFIX = "__tool__:"

_TOOL_STATUS_LABELS: dict[str, str] = {
    "thinking": "🤔 思考中...",
    "searching": "🔍 搜索中...",
    "reading": "📖 阅读中...",
    "generating": "✍️ 生成中...",
    "executing": "⚙️ 执行中...",
    "summarizing": "📝 总结中...",
}


def get_tool_status_label(tool_status: str) -> str:
    return _TOOL_STATUS_LABELS.get(tool_status, f"... {tool_status} ...")


class DraftStreamLoop:
    """流式草稿消息循环引擎。

    ``update()`` 方法要求调用方每次传入**完整的累积文本**
    （而非增量 chunk），因为内部 ``StreamingDirectiveParser``
    需要对全文做指令扫描。传入增量文本会导致未发送内容静默丢失。
    """

    def __init__(self, send_or_edit: SendOrEditFn, throttle_ms: int = 1000, max_retries: int = 3):
        self._send_or_edit = send_or_edit
        self._throttle_ms = throttle_ms
        self._active_session: DraftStreamSession | None = None
        self._update_queue: asyncio.Queue[str | None] | None = None
        self._loop_task: asyncio.Task[None] | None = None
        self._lock = asyncio.Lock()

        self._last_sent_at: float = 0.0
        self._pending_text: str = ""
        self._in_flight: asyncio.Task[str | None] | None = None
        self._timer: asyncio.Task[None] | None = None
        self._flush_task: asyncio.Task[None] | None = None

        self._directive_parser = StreamingDirectiveParser()
        self._on_directive: DirectiveHandler | None = None

        self._on_progress: ProgressCallback | None = None

        self._send_backoff = ErrorBackoff(
            config=BackoffConfig(
                base_delay=0.5,
                max_delay=5.0,
                exponent=2.0,
                jitter=True,
                jitter_factor=0.1,
                max_retries=max_retries,
            ),
        )

    def set_directive_handler(self, handler: DirectiveHandler) -> None:
        self._on_directive = handler

    def set_progress_callback(self, callback: ProgressCallback) -> None:
        self._on_progress = callback

    async def start(self, session: DraftStreamSession) -> None:
        async with self._lock:
            if self._active_session is not None:
                await self._abort_current()
            self._active_session = session
            session.state = DraftStreamState.STREAMING
            self._update_queue = asyncio.Queue()
            self._directive_parser.reset()
            self._loop_task = asyncio.create_task(self._run_loop(session))

    async def update(self, full_text: str, tool_status: str | None = None) -> None:
        """将**完整累积文本**送入流式处理循环。

        调用方必须传入当前已累积的全部文本内容，而非增量 chunk。
        内部 ``StreamingDirectiveParser`` 会对全文做指令解析与去重。
        """
        if self._update_queue is None:
            return
        await self._update_queue.put(full_text)
        if tool_status is not None:
            await self._update_queue.put(f"{_TOOL_STATUS_PREFIX}{tool_status}")

    async def finalize(self) -> None:
        if self._update_queue is not None:
            await self._update_queue.put(None)
        if self._loop_task is not None:
            try:
                await asyncio.wait_for(self._loop_task, timeout=TIMEOUT.stream_finalize)
            except TimeoutError:
                logger.warning("DraftStreamLoop finalize timed out after 30s, forcing cancel")
                self._loop_task.cancel()
                try:
                    await self._loop_task
                except asyncio.CancelledError:
                    pass
            finally:
                self._loop_task = None
        if self._active_session is not None and self._active_session.state == DraftStreamState.FINISHED:
            self._active_session = None

    async def abort(self) -> None:
        async with self._lock:
            await self._abort_current()

    async def flush(self) -> None:
        self._cancel_timer()
        iterations = 0
        while not self._is_stopped and iterations < _MAX_FLUSH_ITERATIONS:
            iterations += 1
            if self._in_flight and not self._in_flight.done():
                await self._in_flight
                continue

            text = self._pending_text
            if not text.strip():
                self._pending_text = ""
                return

            self._pending_text = ""
            msg_id = self._active_session.message_id if self._active_session else None
            key = f"flush:{id(self)}"

            async def _do_flush() -> str | None:
                return await self._send_or_edit(text, msg_id)

            async def _retry_flush() -> str | None:
                return await self._send_backoff.execute(key, _do_flush)

            current = asyncio.create_task(_retry_flush())

            def _clear_in_flight(t: asyncio.Task[str | None]) -> None:
                if self._in_flight is t:
                    self._in_flight = None

            current.add_done_callback(_clear_in_flight)
            self._in_flight = current

            try:
                sent = await current
                if sent is None:
                    self._pending_text = text
                    return
                self._last_sent_at = time.monotonic()
                if not self._pending_text:
                    return
            except Exception:
                logger.exception(
                    "DraftStreamLoop flush failed after %d retries",
                    self._send_backoff.config.max_retries,
                )
                self._pending_text = text
                return

        if iterations >= _MAX_FLUSH_ITERATIONS:
            logger.warning("DraftStreamLoop.flush() exceeded max iterations (%s), breaking", _MAX_FLUSH_ITERATIONS)

    async def update_text(self, text: str) -> None:
        if self._is_stopped:
            return
        async with self._lock:
            self._pending_text = text
            if self._in_flight and not self._in_flight.done():
                self._schedule()
                return
            if self._timer is None and (time.monotonic() - self._last_sent_at) >= (self._throttle_ms / 1000):
                try:
                    self._flush_task = asyncio.create_task(self.flush())
                except RuntimeError:
                    return
                self._flush_task.add_done_callback(self._on_flush_done)
                return
            self._schedule()

    def stop(self) -> None:
        self._pending_text = ""
        self._cancel_timer()

    def reset_pending(self) -> None:
        self._pending_text = ""

    def reset_throttle_window(self) -> None:
        self._last_sent_at = 0.0
        self._cancel_timer()

    async def wait_for_in_flight(self) -> None:
        if self._in_flight and not self._in_flight.done():
            await self._in_flight

    @property
    def _is_stopped(self) -> bool:
        if self._active_session is None:
            return True
        return self._active_session.state not in (DraftStreamState.STREAMING, DraftStreamState.FINALIZING)

    async def _abort_current(self) -> None:
        if self._loop_task and not self._loop_task.done():
            self._loop_task.cancel()
            try:
                await self._loop_task
            except asyncio.CancelledError:
                pass
            self._loop_task = None
        if self._active_session:
            self._active_session.state = DraftStreamState.ABORTED
            self._active_session = None
        self._update_queue = None
        self._cancel_timer()
        self._pending_text = ""
        self._directive_parser.reset()

    async def _run_loop(self, session: DraftStreamSession) -> None:
        last_update = time.monotonic()
        pending_text = ""

        try:
            while session.state == DraftStreamState.STREAMING:
                try:
                    item = await asyncio.wait_for(
                        self._update_queue.get(),
                        timeout=TIMEOUT.stream_queue_poll,
                    )
                except TimeoutError:
                    now = time.monotonic()
                    if pending_text and (now - last_update) >= (self._throttle_ms / 1000):
                        last_update = now
                        await self._send_update(session, pending_text)
                        pending_text = ""
                        session.last_update_ts = time.monotonic()
                    continue

                if item is None:
                    if pending_text:
                        await self._send_update(session, pending_text)
                    session.state = DraftStreamState.FINISHED
                    return

                if item.startswith(_TOOL_STATUS_PREFIX):
                    status = item[len(_TOOL_STATUS_PREFIX):]
                    session.tool_status = status
                    if self._on_progress is not None:
                        try:
                            await self._on_progress(status)
                        except Exception:
                            logger.exception("Progress callback failed for status %s", status)
                    if pending_text:
                        await self._send_update(session, pending_text)
                        pending_text = ""
                    continue

                clean_text, directives = self._directive_parser.feed(item)
                await self._dispatch_directives(directives)
                pending_text = clean_text
                session.accumulated_text = clean_text

                now = time.monotonic()
                elapsed_ms = int((now - last_update) * 1000)
                delta = len(pending_text) - len(session.last_sent_text)

                is_natural_boundary = bool(pending_text and pending_text[-1] in "\n。！？!?；;：:")

                should_send = (
                    delta >= session.significant_delta_chars
                    or is_natural_boundary
                    or (delta > 0 and elapsed_ms >= session.throttle_ms)
                )

                if should_send and pending_text:
                    last_update = now
                    await self._send_update(session, pending_text)
                    pending_text = ""
                    session.last_update_ts = time.monotonic()
        except asyncio.CancelledError:
            session.state = DraftStreamState.ABORTED
            raise

    async def _send_update(self, session: DraftStreamSession, text: str) -> None:
        key = f"send_update:{id(self)}"

        async def _do_send() -> None:
            new_id = await self._send_or_edit(text, session.message_id)
            if new_id:
                session.message_id = new_id
            session.last_sent_text = text

        try:
            await self._send_backoff.execute(key, _do_send)
        except Exception:
            logger.exception(
                "DraftStreamLoop send_or_edit failed after %d retries",
                self._send_backoff.config.max_retries,
            )

    async def _dispatch_directives(self, directives: list[StreamingDirective]) -> None:
        if not directives or self._on_directive is None:
            return
        for directive in directives:
            try:
                await self._on_directive(directive)
            except Exception:
                logger.exception("Directive handler failed for %s", directive.type)

    def _schedule(self) -> None:
        if self._timer is not None:
            return
        delay = max(0, (self._throttle_ms / 1000) - (time.monotonic() - self._last_sent_at))

        async def _delayed_flush():
            await asyncio.sleep(delay)
            await self.flush()

        try:
            self._timer = asyncio.create_task(_delayed_flush())
        except RuntimeError:
            return
        self._timer.add_done_callback(self._on_timer_done)

    def _cancel_timer(self) -> None:
        if self._timer is not None and not self._timer.done():
            self._timer.cancel()
        self._timer = None

    def _on_flush_done(self, task: asyncio.Task) -> None:
        if self._flush_task is task:
            self._flush_task = None
        try:
            task.result()
        except Exception:
            logger.exception("DraftStreamLoop background flush failed")

    def _on_timer_done(self, task: asyncio.Task) -> None:
        if self._timer is task:
            self._timer = None
        try:
            task.result()
        except Exception:
            logger.exception("DraftStreamLoop timer flush failed")


class FinalizableDraftStreamControls:
    def __init__(
        self,
        stream_loop: DraftStreamLoop,
        session: DraftStreamSession,
        target_id: str,
        account: dict | None = None,
    ):
        self._stream_loop = stream_loop
        self._session = session
        self._target_id = target_id
        self._account = account or {}
        self._started = False

    async def start(self) -> None:
        if self._started:
            return
        await self._stream_loop.start(self._session)
        self._started = True

    async def update(self, full_text: str, tool_status: str | None = None) -> None:
        if not self._started:
            await self.start()
        await self._stream_loop.update(full_text, tool_status)

    async def finalize(self) -> None:
        if self._started:
            await self._stream_loop.finalize()
            self._started = False

    async def abort(self) -> None:
        if self._started:
            await self._stream_loop.abort()
            self._started = False

    def set_progress_callback(self, callback: ProgressCallback) -> None:
        self._stream_loop.set_progress_callback(callback)

    @property
    def session(self) -> DraftStreamSession:
        return self._session

    async def update_text(self, text: str) -> None:
        if self._session.stopped or self._session.final:
            return
        await self.start()
        await self._stream_loop.update_text(text)

    async def stop(self) -> None:
        self._session.final = True
        await self._stream_loop.flush()

    async def seal(self) -> None:
        self._session.final = True
        self._stream_loop.stop()
        await self._stream_loop.wait_for_in_flight()

    def discard_pending(self) -> None:
        self._session.stopped = True
        self._stream_loop.stop()