"""Stdout capture and SSE streaming utilities."""

import sys
import asyncio
import json
from typing import Any, AsyncGenerator, Optional, TextIO
from contextlib import contextmanager

from ._internal.http import notify_run_complete


class StreamingWriter:
    """Custom stdout writer that captures output to an async queue."""

    def __init__(self, queue: asyncio.Queue, original_stdout: TextIO):
        self.queue = queue
        self.original_stdout = original_stdout
        self._buffer = ""

    def write(self, text: str) -> int:
        if text and text.strip():
            # Put in queue for streaming
            try:
                self.queue.put_nowait(text.rstrip("\n"))
            except asyncio.QueueFull:
                pass
        # Also write to original stdout for local debugging
        return self.original_stdout.write(text)

    def flush(self) -> None:
        self.original_stdout.flush()

    def fileno(self) -> int:
        return self.original_stdout.fileno()

    def isatty(self) -> bool:
        return self.original_stdout.isatty()


@contextmanager
def capture_stdout(queue: asyncio.Queue):
    """Context manager to capture stdout to a queue."""
    original_stdout = sys.stdout
    streaming_writer = StreamingWriter(queue, original_stdout)
    sys.stdout = streaming_writer  # type: ignore
    try:
        yield streaming_writer
    finally:
        sys.stdout = original_stdout


def format_sse_event(event_type: str, data: Any) -> str:
    """Format data as an SSE event."""
    payload = json.dumps({"type": event_type, "data": data})
    return f"data: {payload}\n\n"


async def create_sse_stream(
    func,
    *args,
    run_id: Optional[str] = None,
    **kwargs,
) -> AsyncGenerator[str, None]:
    """
    Create an SSE stream that captures stdout and streams it.

    Args:
        func: Async function to execute
        *args: Positional arguments for func
        **kwargs: Keyword arguments for func

    Yields:
        SSE-formatted events
    """
    queue: asyncio.Queue = asyncio.Queue(maxsize=1000)

    # Create task for the user function
    async def run_with_capture():
        with capture_stdout(queue):
            return await func(*args, **kwargs)

    task = asyncio.create_task(run_with_capture())

    # Stream logs while function runs
    try:
        while not task.done():
            try:
                # Check for new output with timeout
                output = await asyncio.wait_for(queue.get(), timeout=0.1)
                yield format_sse_event("log", {"msg": output})
            except asyncio.TimeoutError:
                # No output yet, continue waiting
                continue

        # Get final result
        result = await task

        # Drain any remaining queue items
        while not queue.empty():
            try:
                output = queue.get_nowait()
                yield format_sse_event("log", {"msg": output})
            except asyncio.QueueEmpty:
                break

        # Send result
        yield format_sse_event("result", result)

        # Notify Worker that the SSE stream completed successfully
        if run_id:
            notify_run_complete(run_id, 200)

    except Exception as e:
        # Send error
        yield format_sse_event("error", {"msg": str(e)})
        # Notify Worker that the SSE stream failed
        if run_id:
            notify_run_complete(run_id, 500)
        raise
