"""Output formatter for deep-agent stream events."""

from __future__ import annotations

import re
from collections.abc import AsyncGenerator
from typing import Any

from app.schemas import WsStreamEnd, WsStreamStart, WsStreamText

# Matches <canvas kind="...">...</canvas> blocks (single-line or multiline).
_CANVAS_BLOCK_RE = re.compile(
    r'<canvas\s+kind=["\']([^"\']+)["\']>(.*?)</canvas>',
    re.DOTALL | re.IGNORECASE,
)


def extract_canvas_block(text: str) -> tuple[str, str | None, str | None]:
    """Strip the first <canvas kind="...">...</canvas> block from *text*.

    Returns ``(visible_text, canvas_content, canvas_kind)``.
    ``canvas_content`` and ``canvas_kind`` are ``None`` when no block is found.
    """
    match = _CANVAS_BLOCK_RE.search(text)
    if not match:
        return text, None, None

    canvas_kind = match.group(1).strip()
    canvas_content = match.group(2).strip()
    visible = text[: match.start()] + text[match.end() :]
    visible = visible.strip()
    return visible, canvas_content, canvas_kind

WsFrame = WsStreamStart | WsStreamText | WsStreamEnd


class StreamFormatter:
    """Convert `(event_type, data)` stream events into websocket frame models."""

    def __init__(self, request_id: str) -> None:
        self.request_id = request_id

    async def format(
        self,
        event_stream: AsyncGenerator[tuple[str, Any], None],
    ) -> AsyncGenerator[WsFrame, None]:
        started = False

        async for event_type, data in event_stream:
            if event_type != "token":
                continue

            if not started:
                yield WsStreamStart(request_id=self.request_id)
                started = True

            text = str(data or "")
            if text:
                yield WsStreamText(request_id=self.request_id, chunk=text)

        if not started:
            yield WsStreamStart(request_id=self.request_id)
        yield WsStreamEnd(request_id=self.request_id)