"""
RCABench API

RCABench - A comprehensive root cause analysis benchmarking platform for microservices

The version of the OpenAPI document: 1.1.51
Contact: team@rcabench.com
Generated by OpenAPI Generator (https://rcabench.openapi-generator.tech)

Do not edit the class manually.
"""  # noqa: E501

from typing import Annotated, Any

from pydantic import Field, StrictFloat, StrictInt, StrictStr, validate_call

from rcabench.openapi.api_client import ApiClient, RequestSerialized
from rcabench.openapi.rest import RESTResponse


class TracesApi:
    """NOTE: This class is auto generated by OpenAPI Generator
    Ref: https://rcabench.openapi-generator.tech

    Do not edit the class manually.
    """

    def __init__(self, api_client=None) -> None:
        if api_client is None:
            api_client = ApiClient.get_default()
        self.api_client = api_client

    @staticmethod
    def bytes_to_generator(data: bytes, chunk_size: int | None = None):
        if not chunk_size:
            yield data
        else:
            for i in range(0, len(data), chunk_size):
                yield data[i : i + chunk_size]

    @validate_call
    def get_trace_events(
        self,
        trace_id: Annotated[StrictStr, Field(description="Trace ID")],
        last_id: Annotated[StrictStr | None, Field(description="Last event ID received")] = None,
        _request_timeout: None
        | Annotated[StrictFloat, Field(gt=0)]
        | tuple[Annotated[StrictStr, Field(gt=0)], Annotated[StrictFloat, Field(gt=0)]] = None,
        _request_auth: dict[StrictStr, Any] | None = None,
        _content_type: StrictStr | None = None,
        _headers: dict[StrictStr, Any] | None = None,
        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
    ):
        """Stream trace events in real-time (SSE Stream using sseclient)

        Returns an SSEClient iterator that yields Server-Sent Events.
        This method requires the 'sseclient-py' package to be installed.

        :param trace_id: Trace ID (required)
        :type trace_id: str
        :param last_id: Last event ID received
        :type last_id: str
        :param _request_timeout: timeout setting for this request
        :type _request_timeout: int, tuple(int, int), optional
        :return: SSEClient event iterator with attributes: event, id, data
        """  # noqa: E501

        try:
            import sseclient
        except ImportError:
            raise ImportError(
                "sseclient-py is required for SSE streaming. "
                "Install with: pip install sseclient-py or uv add sseclient-py"
            )

        _param = self._get_trace_events_serialize(
            trace_id=trace_id,
            last_id=last_id,
            _request_auth=_request_auth,
            _content_type=None,
            _headers=_headers,
            _host_index=_host_index,
        )

        _response_types_map: dict[str, str | None] = {
            "200": "str",
            "400": "GenericResponseAny",
            "401": "GenericResponseAny",
            "403": "GenericResponseAny",
            "500": "GenericResponseAny",
        }
        for code in _response_types_map:
            if code.startswith("2"):
                _response_types_map.pop(code)

        response_data = self.api_client.call_api(*_param, _request_timeout=_request_timeout)

        assert response_data is not None and isinstance(response_data, RESTResponse), (
            "Response data must be of type RESTResponse"
        )
        if response_data.status >= 400:
            self.api_client.response_deserialize(
                response_data=response_data,
                response_types_map=_response_types_map,
            )
            return None

        return sseclient.SSEClient(self.bytes_to_generator(response_data.read(), chunk_size=1024))

    def _get_trace_events_serialize(
        self,
        trace_id,
        last_id,
        _request_auth,
        _content_type,
        _headers,
        _host_index,
    ) -> RequestSerialized:
        _host = None

        _collection_formats: dict[str, str] = {}

        _path_params: dict[str, str] = {}
        _query_params: list[tuple[str, str]] = []
        _header_params: dict[str, str | None] = _headers or {}
        _form_params: list[tuple[str, str]] = []
        _files: dict[str, str | bytes | list[str] | list[bytes] | list[tuple[str, bytes]]] = {}
        _body_params: bytes | None = None

        # process the path parameters
        if trace_id is not None:
            _path_params["trace_id"] = trace_id
        # process the query parameters
        if last_id is not None:
            _query_params.append(("last_id", last_id))

        # process the header parameters
        # process the form parameters
        # process the body parameter

        # set the HTTP header `Accept`
        if "Accept" not in _header_params:
            _header_params["Accept"] = self.api_client.select_header_accept(["application/json"])

        # authentication setting
        _auth_settings: list[str] = ["BearerAuth"]

        return self.api_client.param_serialize(
            method="GET",
            resource_path="/api/v2/traces/{trace_id}/stream",
            path_params=_path_params,
            query_params=_query_params,
            header_params=_header_params,
            body=_body_params,
            post_params=_form_params,
            files=_files,
            auth_settings=_auth_settings,
            collection_formats=_collection_formats,
            _host=_host,
            _request_auth=_request_auth,
        )
