"""In-memory generation event buffer for SSE fan-out.

Each active generation gets a GenerationBuffer keyed by conversation_id.
SSE clients tail the buffer and can reconnect at any point without data loss.
"""

import asyncio
import json
import logging
import time
from dataclasses import dataclass, field
from enum import Enum

logger = logging.getLogger(__name__)


class GenerationState(str, Enum):
    RUNNING = "running"
    COMPLETED = "completed"
    ERRORED = "errored"


@dataclass
class SSEEvent:
    index: int
    event_type: str  # "context", "chunk", "done", "error"
    data: dict


@dataclass
class GenerationBuffer:
    conversation_id: int
    assistant_message_id: int
    state: GenerationState = GenerationState.RUNNING
    events: list[SSEEvent] = field(default_factory=list)
    content_so_far: str = ""
    finished_at: float | None = None
    cancel_event: asyncio.Event = field(default_factory=asyncio.Event)
    _notify: asyncio.Event = field(default_factory=asyncio.Event)
    # Tool confirmation — set when a write tool is awaiting user approval
    confirmation_future: asyncio.Future | None = None
    pending_tool: dict | None = None

    def append_event(self, event_type: str, data: dict) -> SSEEvent:
        event = SSEEvent(index=len(self.events), event_type=event_type, data=data)
        self.events.append(event)
        # Wake all waiting SSE clients, then reset for next wait
        old = self._notify
        self._notify = asyncio.Event()
        old.set()
        return event

    async def wait_for_event(self, after_index: int, timeout: float = 30.0) -> bool:
        """Wait until there are events beyond after_index. Returns True if new events exist."""
        if after_index + 1 < len(self.events):
            return True
        try:
            await asyncio.wait_for(self._notify.wait(), timeout=timeout)
            return True
        except asyncio.TimeoutError:
            return False

    def events_after(self, index: int) -> list[SSEEvent]:
        """Return events with index > given index (for replay on reconnect)."""
        start = index + 1
        if start >= len(self.events):
            return []
        return self.events[start:]

    @staticmethod
    def format_sse(event: SSEEvent) -> str:
        return f"id: {event.index}\nevent: {event.event_type}\ndata: {json.dumps(event.data)}\n\n"


# Module-level singleton registry
_buffers: dict[int | str, GenerationBuffer] = {}

_cleanup_task: asyncio.Task | None = None
_GRACE_PERIOD = 60.0  # seconds to keep completed buffers


def create_buffer(conv_id: int, msg_id: int) -> GenerationBuffer:
    existing = _buffers.get(conv_id)
    if existing and existing.state == GenerationState.RUNNING:
        raise RuntimeError(f"Generation already running for conversation {conv_id}")
    buf = GenerationBuffer(conversation_id=conv_id, assistant_message_id=msg_id)
    _buffers[conv_id] = buf
    return buf


def get_buffer(conv_id: int) -> GenerationBuffer | None:
    return _buffers.get(conv_id)


def remove_buffer(conv_id: int) -> None:
    _buffers.pop(conv_id, None)


def create_assist_buffer(user_id: int) -> GenerationBuffer:
    key = f"assist:{user_id}"
    existing = _buffers.get(key)
    if existing and existing.state == GenerationState.RUNNING:
        # Orphan the old buffer — the background task holds a direct reference
        # and will complete against it harmlessly. A new request always wins.
        logger.warning("Assist generation still running for user %d; orphaning old buffer", user_id)
    buf = GenerationBuffer(conversation_id=0, assistant_message_id=0)
    _buffers[key] = buf
    return buf


def get_assist_buffer(user_id: int) -> GenerationBuffer | None:
    return _buffers.get(f"assist:{user_id}")


def remove_assist_buffer(user_id: int) -> None:
    _buffers.pop(f"assist:{user_id}", None)


async def _cleanup_loop() -> None:
    """Remove completed/errored buffers after grace period."""
    while True:
        await asyncio.sleep(15)
        now = time.monotonic()
        to_remove = [
            key for key, buf in _buffers.items()
            if buf.state != GenerationState.RUNNING
            and buf.finished_at is not None
            and now - buf.finished_at > _GRACE_PERIOD
        ]
        for key in to_remove:
            _buffers.pop(key, None)
            logger.debug("Cleaned up generation buffer for key %s", key)


def start_cleanup_loop() -> None:
    global _cleanup_task
    if _cleanup_task is None or _cleanup_task.done():
        _cleanup_task = asyncio.create_task(_cleanup_loop())