from __future__ import annotations

from typing import Callable, Iterator

from ..models import Usage
from .replay import InMemoryReplayStore, ReplayStore, StreamEvent


class ResumeConflict(ValueError):
    def __init__(self, *, last_good_event_id: int) -> None:
        super().__init__("resume_required")
        self.last_good_event_id = int(last_good_event_id)


class StreamCoordinator:
    """Turns provider responses into streaming events with replay support."""

    def __init__(self, replay_store: ReplayStore | None = None, token_size: int = 20) -> None:
        self._replay = replay_store or InMemoryReplayStore()
        self._token_size = max(1, token_size)
        self._next_event_id: dict[str, int] = {}

    def stream_text(
        self,
        request_id: str,
        text: str,
        *,
        max_completion_tokens: int | None = None,
        is_cancelled: Callable[[str], bool] | None = None,
    ) -> Iterator[StreamEvent]:
        words = text.split()
        parts: list[str] = []
        chunk: list[str] = []
        for word in words:
            chunk.append(word)
            if len(chunk) >= self._token_size:
                parts.append(" ".join(chunk))
                chunk = []
        if chunk:
            parts.append(" ".join(chunk))

        event_id = self._next_event_id.get(request_id, 0)
        emitted_completion_tokens = 0
        for part in parts:
            if is_cancelled and is_cancelled(request_id):
                break
            event_id += 1
            prompt_tokens = max(0, len(part.split()) // 2)
            completion_tokens = max(1, len(part.split()))
            if (
                max_completion_tokens is not None
                and emitted_completion_tokens + completion_tokens > max_completion_tokens
            ):
                remaining = max(0, max_completion_tokens - emitted_completion_tokens)
                if remaining <= 0:
                    break
                words_subset = part.split()[:remaining]
                part = " ".join(words_subset)
                completion_tokens = len(words_subset)
            usage_delta = Usage(
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                total_tokens=prompt_tokens + completion_tokens,
            )
            event = StreamEvent(event_id=event_id, data=part, usage_delta=usage_delta)
            self._replay.append(request_id, event)
            emitted_completion_tokens += completion_tokens
            yield event
        self._next_event_id[request_id] = event_id

    def resume(self, request_id: str, last_event_id: int) -> list[StreamEvent]:
        return self._replay.replay_from(request_id, last_event_id)

    def resume_checked(self, request_id: str, after_event_id: int) -> list[StreamEvent]:
        after_event_id = int(after_event_id)
        if after_event_id < 0:
            raise ValueError("after_event_id must be >= 0")

        all_events = self.resume(request_id, 0)
        min_event_id = min((ev.event_id for ev in all_events), default=0)
        max_event_id = max((ev.event_id for ev in all_events), default=0)

        if after_event_id > max_event_id:
            raise ResumeConflict(last_good_event_id=max_event_id)

        if min_event_id and after_event_id < min_event_id - 1:
            raise ResumeConflict(last_good_event_id=min_event_id - 1)

        return [ev for ev in all_events if ev.event_id > after_event_id]


__all__ = ["ResumeConflict", "StreamCoordinator"]
