from __future__ import annotations

import asyncio
import logging
import time
from dataclasses import dataclass, field
from enum import Enum, auto
from collections.abc import AsyncGenerator

logger = logging.getLogger(__name__)


class StreamEventType(Enum):
    FLUSH = auto()
    COMPLETE = auto()
    ERROR = auto()


class FlushStrategy(Enum):
    PER_CHAR = auto()
    INTERVAL = auto()
    BACKLOG = auto()


@dataclass
class StreamEvent:
    event_type: StreamEventType
    content: str = ""
    full_content: str = ""
    error: str | None = None


@dataclass
class StreamCheckpoint:
    chat_id: str
    msg_id: str = ""
    content: str = ""
    sent_seq: int = 0
    timestamp: float = field(default_factory=time.time)

    @property
    def is_stale(self) -> bool:
        return time.time() - self.timestamp > 600


class FlushController:
    strategy: FlushStrategy = FlushStrategy.BACKLOG
    flush_interval: float = 0.6
    backlog_threshold: int = 30
    max_retries: int = 3
    retry_delay: float = 1.0
    seq_counter: int = 0
    _buffer: list[str] = field(default_factory=list, repr=False)
    _last_flush: float = 0.0
    _accumulated: str = ""

    def reset(self) -> None:
        self._buffer = []
        self._last_flush = 0.0
        self._accumulated = ""
        self.seq_counter = 0

    def feed(self, chunk: str) -> list[str]:
        self._buffer.append(chunk)
        self._accumulated += chunk

        if self.strategy == FlushStrategy.PER_CHAR:
            return self._flush_all()

        if self.strategy == FlushStrategy.BACKLOG and len(self._accumulated) >= self.backlog_threshold:
            return self._flush_all()

        now = time.monotonic()
        if self.strategy == FlushStrategy.INTERVAL and now - self._last_flush >= self.flush_interval:
            return self._flush_all()

        return []

    def flush_remaining(self) -> tuple[list[str], str]:
        if self.strategy == FlushStrategy.PER_CHAR:
            return [], self._accumulated
        batches = self._flush_all() if self._buffer else []
        return batches, self._accumulated

    def _flush_all(self) -> list[str]:
        if not self._buffer:
            return []
        batches = list(self._buffer)
        self._buffer = []
        self._accumulated = ""
        self._last_flush = time.monotonic()
        return batches


class C2CStreamingController:
    def __init__(
        self,
        send_message_fn,
        retry_delay: float = 1.0,
        max_retries: int = 3,
        flush_interval: float = 0.6,
    ) -> None:
        self._send_message_fn = send_message_fn
        self._retry_delay = retry_delay
        self._max_retries = max_retries
        self._flush_interval = flush_interval
        self._flush_controller = FlushController(strategy=FlushStrategy.BACKLOG, flush_interval=flush_interval)
        self._any_chunk_delivered: bool = False
        self._static_fallback_msg: str = ""
        self._active_checkpoints: dict[str, StreamCheckpoint] = {}

    @property
    def flush_controller(self) -> FlushController:
        return self._flush_controller

    async def stream(
        self,
        chat_id: str,
        msg_id: str,
        content_generator: AsyncGenerator[str, None],
        event_id: str = "",
    ) -> bool:
        self._flush_controller.reset()
        self._any_chunk_delivered = False
        self._static_fallback_msg = ""

        checkpoint = StreamCheckpoint(chat_id=chat_id, msg_id=msg_id)
        self._active_checkpoints[chat_id] = checkpoint

        collected = ""
        in_media_interrupt = False
        media_buffer = ""

        try:
            async for chunk in content_generator:
                if not chunk:
                    continue

                media_splits = self._split_on_media_tags(chunk)
                for segment, is_media in media_splits:
                    if is_media:
                        in_media_interrupt = True
                        media_buffer += segment
                    else:
                        if in_media_interrupt and media_buffer:
                            checkpoint.content = collected
                            await self._flush_and_end_stream(chat_id, msg_id, checkpoint, event_id)
                            await self._send_media_interruption(chat_id, media_buffer)
                            media_buffer = ""
                            in_media_interrupt = False

                            new_msg_id = f"{msg_id}_resume_{checkpoint.sent_seq}"
                            checkpoint.msg_id = new_msg_id
                            self._flush_controller.reset()

                        collected += segment
                        batches = self._flush_controller.feed(segment)
                        for batch in batches:
                            self._flush_controller.seq_counter += 1
                            success = await self._send_stream_chunk(
                                chat_id,
                                checkpoint.msg_id or msg_id,
                                batch,
                                self._flush_controller.seq_counter,
                                event_id,
                            )
                            if success:
                                self._any_chunk_delivered = True
                            else:
                                logger.warning(
                                    "C2CStreaming: flush failed for seq %d", self._flush_controller.seq_counter
                                )

            self._static_fallback_msg = collected

            flush_batches, full = self._flush_controller.flush_remaining()
            for batch in flush_batches:
                self._flush_controller.seq_counter += 1
                success = await self._send_stream_chunk(
                    chat_id,
                    checkpoint.msg_id or msg_id,
                    batch,
                    self._flush_controller.seq_counter,
                    event_id,
                )
                if success:
                    self._any_chunk_delivered = True

            self._flush_controller.seq_counter += 1
            success = await self._send_stream_chunk(
                chat_id,
                checkpoint.msg_id or msg_id,
                "",
                self._flush_controller.seq_counter,
                event_id,
                is_end=True,
            )
            self._active_checkpoints.pop(chat_id, None)
            return success

        except Exception as e:
            logger.exception("C2CStreaming: stream error for %s", chat_id)
            self._active_checkpoints.pop(chat_id, None)
            return await self._cancel_stream(chat_id, msg_id, event_id, str(e))

    async def _send_stream_chunk(
        self,
        chat_id: str,
        msg_id: str,
        content: str,
        msg_seq: int,
        event_id: str = "",
        is_end: bool = False,
    ) -> bool:
        for attempt in range(self._max_retries):
            try:
                payload = {
                    "content": content,
                    "msg_type": 0,
                    "msg_id": msg_id,
                    "msg_seq": msg_seq,
                    "stream": {"state": 2 if is_end else 1},
                }
                if event_id:
                    payload["event_id"] = event_id

                response = await self._send_message_fn(chat_id, payload)
                if response is not None:
                    cp = self._active_checkpoints.get(chat_id)
                    if cp:
                        cp.sent_seq = msg_seq
                        cp.timestamp = time.time()
                return response is not None
            except Exception:
                if attempt < self._max_retries - 1:
                    await asyncio.sleep(self._retry_delay)
                else:
                    logger.exception("C2CStreaming: send chunk failed after %d retries", self._max_retries)
        return False

    @staticmethod
    def _split_on_media_tags(text: str) -> list[tuple[str, bool]]:
        import re

        from .media_tags import _MEDIA_TAG_RE, _IMG_TAG_RE, _VIDEO_TAG_RE

        combined_re = re.compile(f"({_MEDIA_TAG_RE.pattern}|{_IMG_TAG_RE.pattern}|{_VIDEO_TAG_RE.pattern})")

        results: list[tuple[str, bool]] = []
        last_end = 0
        for match in combined_re.finditer(text):
            if match.start() > last_end:
                results.append((text[last_end : match.start()], False))
            results.append((match.group(0), True))
            last_end = match.end()

        if last_end < len(text):
            results.append((text[last_end:], False))

        return results

    async def _flush_and_end_stream(
        self, chat_id: str, msg_id: str, checkpoint: StreamCheckpoint, event_id: str
    ) -> None:
        flush_batches, _ = self._flush_controller.flush_remaining()
        for batch in flush_batches:
            self._flush_controller.seq_counter += 1
            await self._send_stream_chunk(
                chat_id,
                msg_id,
                batch,
                self._flush_controller.seq_counter,
                event_id,
            )

        self._flush_controller.seq_counter += 1
        await self._send_stream_chunk(
            chat_id,
            msg_id,
            "",
            self._flush_controller.seq_counter,
            event_id,
            is_end=True,
        )

    async def _send_media_interruption(self, chat_id: str, media_tags: str) -> None:
        try:
            payload = {
                "content": media_tags,
                "msg_type": 0,
                "msg_id": "",
                "msg_seq": 0,
            }
            await self._send_message_fn(chat_id, payload)
        except Exception:
            logger.exception("C2CStreaming: media interruption send failed")

    async def _cancel_stream(self, chat_id: str, msg_id: str, event_id: str, reason: str) -> bool:
        delivered_fallback = False
        if self._static_fallback_msg and not self._any_chunk_delivered:
            try:
                static_payload = {
                    "content": self._static_fallback_msg[:2000],
                    "msg_type": 0,
                    "msg_id": msg_id,
                    "msg_seq": 0,
                }
                if event_id:
                    static_payload["event_id"] = event_id
                await self._send_message_fn(chat_id, static_payload)
                delivered_fallback = True
                logger.info("C2CStreaming: delivered static fallback message for %s", chat_id)
            except Exception:
                logger.exception("C2CStreaming: static fallback delivery failed for %s", chat_id)

        if not delivered_fallback:
            try:
                payload = {
                    "content": reason,
                    "msg_type": 0,
                    "msg_id": msg_id,
                    "msg_seq": 0,
                    "stream": {"state": 0},
                }
                if event_id:
                    payload["event_id"] = event_id
                await self._send_message_fn(chat_id, payload)
            except Exception:
                pass
        return False

    def get_checkpoint(self, chat_id: str) -> StreamCheckpoint | None:
        return self._active_checkpoints.get(chat_id)

    def has_pending_stream(self, chat_id: str) -> bool:
        cp = self._active_checkpoints.get(chat_id)
        return cp is not None and not cp.is_stale

    def cleanup_stale_checkpoints(self) -> int:
        stale = [cid for cid, cp in self._active_checkpoints.items() if cp.is_stale]
        for cid in stale:
            self._active_checkpoints.pop(cid, None)
        return len(stale)