from __future__ import annotations

from dataclasses import dataclass
from enum import StrEnum

from yuxi.channel.config.defaults import (
    BLOCK_STREAMING_BREAK,
    BLOCK_STREAMING_CHUNK_BREAK_PREFERENCE,
    BLOCK_STREAMING_CHUNK_MAX_CHARS,
    BLOCK_STREAMING_CHUNK_MIN_CHARS,
    BLOCK_STREAMING_COALESCE_IDLE_MS,
    BLOCK_STREAMING_COALESCE_JOINER,
    BLOCK_STREAMING_COALESCE_MAX_CHARS,
    BLOCK_STREAMING_COALESCE_MIN_CHARS,
    STREAMING_PREVIEW_MIN_INITIAL_CHARS,
    STREAMING_PREVIEW_THROTTLE_MS,
)


class BlockStreamingBreak(StrEnum):
    TEXT_END = "text_end"


class DraftStreamState(StrEnum):
    IDLE = "idle"
    STREAMING = "streaming"
    FINALIZING = "finalizing"
    FINISHED = "finished"
    ABORTED = "aborted"


@dataclass
class StreamingConfig:
    preview_stream_throttle_ms: int = STREAMING_PREVIEW_THROTTLE_MS
    preview_min_initial_chars: int = STREAMING_PREVIEW_MIN_INITIAL_CHARS

    block_streaming_enabled: bool = False
    block_streaming_break: str = BLOCK_STREAMING_BREAK
    block_streaming_chunk_min_chars: int = BLOCK_STREAMING_CHUNK_MIN_CHARS
    block_streaming_chunk_max_chars: int = BLOCK_STREAMING_CHUNK_MAX_CHARS
    block_streaming_chunk_break_preference: str = BLOCK_STREAMING_CHUNK_BREAK_PREFERENCE

    block_streaming_flush_on_paragraph: bool = False

    block_streaming_coalesce_min_chars: int | None = BLOCK_STREAMING_COALESCE_MIN_CHARS
    block_streaming_coalesce_max_chars: int | None = BLOCK_STREAMING_COALESCE_MAX_CHARS
    block_streaming_coalesce_idle_ms: int = BLOCK_STREAMING_COALESCE_IDLE_MS
    block_streaming_coalesce_joiner: str = BLOCK_STREAMING_COALESCE_JOINER

    @classmethod
    def defaults(cls) -> dict:
        return {
            "preview_stream_throttle_ms": STREAMING_PREVIEW_THROTTLE_MS,
            "preview_min_initial_chars": STREAMING_PREVIEW_MIN_INITIAL_CHARS,
            "block_streaming_enabled": False,
            "block_streaming_break": BLOCK_STREAMING_BREAK,
            "block_streaming_chunk_min_chars": BLOCK_STREAMING_CHUNK_MIN_CHARS,
            "block_streaming_chunk_max_chars": BLOCK_STREAMING_CHUNK_MAX_CHARS,
            "block_streaming_chunk_break_preference": BLOCK_STREAMING_CHUNK_BREAK_PREFERENCE,
            "block_streaming_flush_on_paragraph": False,
            "block_streaming_coalesce_min_chars": BLOCK_STREAMING_COALESCE_MIN_CHARS,
            "block_streaming_coalesce_max_chars": BLOCK_STREAMING_COALESCE_MAX_CHARS,
            "block_streaming_coalesce_idle_ms": BLOCK_STREAMING_COALESCE_IDLE_MS,
            "block_streaming_coalesce_joiner": BLOCK_STREAMING_COALESCE_JOINER,
        }

    @classmethod
    def empty(cls) -> StreamingConfig:
        return cls()


@dataclass
class DraftStreamSession:
    message_id: str | None = None
    throttle_ms: int = STREAMING_PREVIEW_THROTTLE_MS
    significant_delta_chars: int = STREAMING_PREVIEW_MIN_INITIAL_CHARS
    state: DraftStreamState = DraftStreamState.IDLE
    accumulated_text: str = ""
    last_sent_text: str = ""
    last_update_ts: float = 0.0
    tool_status: str | None = None

    stopped: bool = False
    final: bool = False

    @property
    def is_active(self) -> bool:
        return self.state in (DraftStreamState.STREAMING, DraftStreamState.FINALIZING)

    @property
    def pending_chars(self) -> str:
        if len(self.accumulated_text) > len(self.last_sent_text):
            return self.accumulated_text[len(self.last_sent_text) :]
        return ""


@dataclass
class BlockReplyChunking:
    min_chars: int = BLOCK_STREAMING_CHUNK_MIN_CHARS
    max_chars: int = BLOCK_STREAMING_CHUNK_MAX_CHARS
    break_preference: str = BLOCK_STREAMING_CHUNK_BREAK_PREFERENCE
    flush_on_paragraph: bool = False


@dataclass
class BlockReplyCoalescing:
    min_chars: int = BLOCK_STREAMING_CHUNK_MIN_CHARS
    max_chars: int = BLOCK_STREAMING_CHUNK_MAX_CHARS
    idle_ms: int = BLOCK_STREAMING_COALESCE_IDLE_MS
    joiner: str = BLOCK_STREAMING_COALESCE_JOINER
    flush_on_enqueue: bool = False