"""Stream wrappers for NDJSON streaming execution responses."""

from __future__ import annotations

from types import TracebackType
from typing import Any, Iterator, Self

import httpx

from baponi._base import (
    build_cancel_url,
    parse_ndjson_line,
    raise_for_status,
)
from baponi._events import ExecutionEvent, ResultEvent
from baponi.exceptions import BaponiError
from baponi.types import CancelResponse, SandboxResult


class ExecutionStream:
    """Sync iterator over streaming execution events.

    Usage::

        with client.execute_stream("print('hello')") as stream:
            for event in stream:
                print(event)
            result = stream.get_final_result()
    """

    def __init__(
        self,
        response: httpx.Response,
        *,
        base_url: str,
        headers: dict[str, str],
        client: httpx.Client,
    ) -> None:
        self._response = response
        self._base_url = base_url
        self._headers = headers
        self._client = client
        self._trace_id: str | None = None
        self._last_seq: int = 0
        self._final_result: SandboxResult | None = None
        self._closed = False
        self._line_iter: Iterator[str] | None = None

    @property
    def trace_id(self) -> str | None:
        return self._trace_id

    @property
    def last_seq(self) -> int:
        return self._last_seq

    def __enter__(self) -> Self:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        self.close()

    def __iter__(self) -> Iterator[ExecutionEvent]:
        return self

    def __next__(self) -> ExecutionEvent:
        if self._line_iter is None:
            self._line_iter = self._response.iter_lines()
        for line in self._line_iter:
            event = parse_ndjson_line(line)
            if event is None:
                continue
            self._last_seq = event.seq
            if hasattr(event, "trace_id") and self._trace_id is None:
                self._trace_id = event.trace_id  # type: ignore[union-attr]
            if isinstance(event, ResultEvent):
                self._final_result = event.result
            return event
        raise StopIteration

    def get_final_result(self) -> SandboxResult:
        """Return the final SandboxResult from the stream.

        Raises BaponiError if the stream did not produce a result event.
        """
        if self._final_result is None:
            raise BaponiError("Stream did not produce a result event")
        return self._final_result

    def until_done(self) -> SandboxResult:
        """Consume all events silently and return the final result."""
        for _ in self:
            pass
        return self.get_final_result()

    def cancel(self) -> CancelResponse:
        """Cancel the running execution."""
        if self._trace_id is None:
            raise BaponiError("Cannot cancel: no trace_id received yet")
        url = build_cancel_url(self._base_url, self._trace_id)
        response = self._client.post(url, headers=self._headers)
        raise_for_status(response)
        return CancelResponse.model_validate(response.json())

    def close(self) -> None:
        """Close the underlying response."""
        if not self._closed:
            self._response.close()
            self._closed = True


class AsyncExecutionStream:
    """Async iterator over streaming execution events.

    Usage::

        async with client.execute_stream("print('hello')") as stream:
            async for event in stream:
                print(event)
            result = stream.get_final_result()
    """

    def __init__(
        self,
        response: httpx.Response,
        *,
        base_url: str,
        headers: dict[str, str],
        client: httpx.AsyncClient,
    ) -> None:
        self._response = response
        self._base_url = base_url
        self._headers = headers
        self._client = client
        self._trace_id: str | None = None
        self._last_seq: int = 0
        self._final_result: SandboxResult | None = None
        self._closed = False
        self._line_iter: Any = None

    @property
    def trace_id(self) -> str | None:
        return self._trace_id

    @property
    def last_seq(self) -> int:
        return self._last_seq

    async def __aenter__(self) -> Self:
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        await self.close()

    def __aiter__(self) -> AsyncExecutionStream:
        return self

    async def __anext__(self) -> ExecutionEvent:
        if self._line_iter is None:
            self._line_iter = self._response.aiter_lines()
        async for line in self._line_iter:
            event = parse_ndjson_line(line)
            if event is None:
                continue
            self._last_seq = event.seq
            if hasattr(event, "trace_id") and self._trace_id is None:
                self._trace_id = event.trace_id  # type: ignore[union-attr]
            if isinstance(event, ResultEvent):
                self._final_result = event.result
            return event
        raise StopAsyncIteration

    def get_final_result(self) -> SandboxResult:
        """Return the final SandboxResult from the stream."""
        if self._final_result is None:
            raise BaponiError("Stream did not produce a result event")
        return self._final_result

    async def until_done(self) -> SandboxResult:
        """Consume all events silently and return the final result."""
        async for _ in self:
            pass
        return self.get_final_result()

    async def cancel(self) -> CancelResponse:
        """Cancel the running execution."""
        if self._trace_id is None:
            raise BaponiError("Cannot cancel: no trace_id received yet")
        url = build_cancel_url(self._base_url, self._trace_id)
        response = await self._client.post(url, headers=self._headers)
        raise_for_status(response)
        return CancelResponse.model_validate(response.json())

    async def close(self) -> None:
        """Close the underlying response."""
        if not self._closed:
            await self._response.aclose()
            self._closed = True
