# This file was auto-generated by Fern from our API Definition.

import typing
from json.decoder import JSONDecodeError

from ..core.api_error import ApiError
from ..core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
from ..core.pagination import AsyncPager, SyncPager
from ..core.pydantic_utilities import parse_obj_as
from ..core.request_options import RequestOptions
from ..core.serialization import convert_and_respect_annotation_metadata
from ..types.query_spans_response import QuerySpansResponse
from ..types.sort_direction import SortDirection
from ..types.subject_type import SubjectType
from ..types.trace_span import TraceSpan
from .types.query_spans_request_filters_item import QuerySpansRequestFiltersItem

# this is used as the default value for optional parameters
OMIT = typing.cast(typing.Any, ...)


class RawTracesClient:
    def __init__(self, *, client_wrapper: SyncClientWrapper):
        self._client_wrapper = client_wrapper

    def query_spans(
        self,
        *,
        start_time: str,
        tracing_project_fqn: str,
        end_time: typing.Optional[str] = OMIT,
        trace_ids: typing.Optional[typing.Sequence[str]] = OMIT,
        span_ids: typing.Optional[typing.Sequence[str]] = OMIT,
        parent_span_ids: typing.Optional[typing.Sequence[str]] = OMIT,
        created_by_subject_types: typing.Optional[typing.Sequence[SubjectType]] = OMIT,
        created_by_subject_slugs: typing.Optional[typing.Sequence[str]] = OMIT,
        application_names: typing.Optional[typing.Sequence[str]] = OMIT,
        limit: typing.Optional[int] = OMIT,
        sort_direction: typing.Optional[SortDirection] = OMIT,
        page_token: typing.Optional[str] = OMIT,
        filters: typing.Optional[typing.Sequence[QuerySpansRequestFiltersItem]] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> SyncPager[TraceSpan, QuerySpansResponse]:
        """
        Parameters
        ----------
        start_time : str
            Start time in ISO 8601 format (e.g., 2025-03-12T00:00:09.872Z)

        tracing_project_fqn : str
            Tracing project FQN (e.g., truefoundry:tracing-project:tfy-default)

        end_time : typing.Optional[str]
            End time in ISO 8601 format (e.g., 2025-03-12T00:10:00.000Z). Defaults to current time if not provided.

        trace_ids : typing.Optional[typing.Sequence[str]]
            Array of trace IDs to filter by

        span_ids : typing.Optional[typing.Sequence[str]]
            Array of span IDs to filter by

        parent_span_ids : typing.Optional[typing.Sequence[str]]
            Array of parent span IDs to filter by

        created_by_subject_types : typing.Optional[typing.Sequence[SubjectType]]
            Array of subject types to filter by

        created_by_subject_slugs : typing.Optional[typing.Sequence[str]]
            Array of subject slugs to filter by

        application_names : typing.Optional[typing.Sequence[str]]
            Array of application names to filter by

        limit : typing.Optional[int]
            The maximum number of spans to return per page. Defaults to 200 if not provided.

        sort_direction : typing.Optional[SortDirection]
            Sort direction for results based on time. Defaults to descending (latest first)

        page_token : typing.Optional[str]
            An opaque string that should be passed as-is from previous response for fetching the next page. Pass `$response.pagination.nextPageToken` from previous response for fetching the next page.

        filters : typing.Optional[typing.Sequence[QuerySpansRequestFiltersItem]]
            Array of filters

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Returns
        -------
        SyncPager[TraceSpan, QuerySpansResponse]
            Returns all the spans matching the query.
        """
        _response = self._client_wrapper.httpx_client.request(
            "api/svc/v1/spans/query",
            method="POST",
            json={
                "startTime": start_time,
                "endTime": end_time,
                "traceIds": trace_ids,
                "spanIds": span_ids,
                "parentSpanIds": parent_span_ids,
                "createdBySubjectTypes": created_by_subject_types,
                "createdBySubjectSlugs": created_by_subject_slugs,
                "applicationNames": application_names,
                "limit": limit,
                "sortDirection": sort_direction,
                "pageToken": page_token,
                "tracingProjectFqn": tracing_project_fqn,
                "filters": convert_and_respect_annotation_metadata(
                    object_=filters, annotation=typing.Sequence[QuerySpansRequestFiltersItem], direction="write"
                ),
            },
            headers={
                "content-type": "application/json",
            },
            request_options=request_options,
            omit=OMIT,
        )
        try:
            if 200 <= _response.status_code < 300:
                _parsed_response = typing.cast(
                    QuerySpansResponse,
                    parse_obj_as(
                        type_=QuerySpansResponse,  # type: ignore
                        object_=_response.json(),
                    ),
                )
                _items = _parsed_response.data
                _has_next = False
                _get_next = None
                if _parsed_response.pagination is not None:
                    _parsed_next = _parsed_response.pagination.next_page_token
                    _has_next = _parsed_next is not None and _parsed_next != ""
                    _get_next = lambda: self.query_spans(
                        start_time=start_time,
                        tracing_project_fqn=tracing_project_fqn,
                        end_time=end_time,
                        trace_ids=trace_ids,
                        span_ids=span_ids,
                        parent_span_ids=parent_span_ids,
                        created_by_subject_types=created_by_subject_types,
                        created_by_subject_slugs=created_by_subject_slugs,
                        application_names=application_names,
                        limit=limit,
                        sort_direction=sort_direction,
                        page_token=_parsed_next,
                        filters=filters,
                        request_options=request_options,
                    )
                return SyncPager(has_next=_has_next, items=_items, get_next=_get_next, response=_parsed_response)
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
        raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)


class AsyncRawTracesClient:
    def __init__(self, *, client_wrapper: AsyncClientWrapper):
        self._client_wrapper = client_wrapper

    async def query_spans(
        self,
        *,
        start_time: str,
        tracing_project_fqn: str,
        end_time: typing.Optional[str] = OMIT,
        trace_ids: typing.Optional[typing.Sequence[str]] = OMIT,
        span_ids: typing.Optional[typing.Sequence[str]] = OMIT,
        parent_span_ids: typing.Optional[typing.Sequence[str]] = OMIT,
        created_by_subject_types: typing.Optional[typing.Sequence[SubjectType]] = OMIT,
        created_by_subject_slugs: typing.Optional[typing.Sequence[str]] = OMIT,
        application_names: typing.Optional[typing.Sequence[str]] = OMIT,
        limit: typing.Optional[int] = OMIT,
        sort_direction: typing.Optional[SortDirection] = OMIT,
        page_token: typing.Optional[str] = OMIT,
        filters: typing.Optional[typing.Sequence[QuerySpansRequestFiltersItem]] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> AsyncPager[TraceSpan, QuerySpansResponse]:
        """
        Parameters
        ----------
        start_time : str
            Start time in ISO 8601 format (e.g., 2025-03-12T00:00:09.872Z)

        tracing_project_fqn : str
            Tracing project FQN (e.g., truefoundry:tracing-project:tfy-default)

        end_time : typing.Optional[str]
            End time in ISO 8601 format (e.g., 2025-03-12T00:10:00.000Z). Defaults to current time if not provided.

        trace_ids : typing.Optional[typing.Sequence[str]]
            Array of trace IDs to filter by

        span_ids : typing.Optional[typing.Sequence[str]]
            Array of span IDs to filter by

        parent_span_ids : typing.Optional[typing.Sequence[str]]
            Array of parent span IDs to filter by

        created_by_subject_types : typing.Optional[typing.Sequence[SubjectType]]
            Array of subject types to filter by

        created_by_subject_slugs : typing.Optional[typing.Sequence[str]]
            Array of subject slugs to filter by

        application_names : typing.Optional[typing.Sequence[str]]
            Array of application names to filter by

        limit : typing.Optional[int]
            The maximum number of spans to return per page. Defaults to 200 if not provided.

        sort_direction : typing.Optional[SortDirection]
            Sort direction for results based on time. Defaults to descending (latest first)

        page_token : typing.Optional[str]
            An opaque string that should be passed as-is from previous response for fetching the next page. Pass `$response.pagination.nextPageToken` from previous response for fetching the next page.

        filters : typing.Optional[typing.Sequence[QuerySpansRequestFiltersItem]]
            Array of filters

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Returns
        -------
        AsyncPager[TraceSpan, QuerySpansResponse]
            Returns all the spans matching the query.
        """
        _response = await self._client_wrapper.httpx_client.request(
            "api/svc/v1/spans/query",
            method="POST",
            json={
                "startTime": start_time,
                "endTime": end_time,
                "traceIds": trace_ids,
                "spanIds": span_ids,
                "parentSpanIds": parent_span_ids,
                "createdBySubjectTypes": created_by_subject_types,
                "createdBySubjectSlugs": created_by_subject_slugs,
                "applicationNames": application_names,
                "limit": limit,
                "sortDirection": sort_direction,
                "pageToken": page_token,
                "tracingProjectFqn": tracing_project_fqn,
                "filters": convert_and_respect_annotation_metadata(
                    object_=filters, annotation=typing.Sequence[QuerySpansRequestFiltersItem], direction="write"
                ),
            },
            headers={
                "content-type": "application/json",
            },
            request_options=request_options,
            omit=OMIT,
        )
        try:
            if 200 <= _response.status_code < 300:
                _parsed_response = typing.cast(
                    QuerySpansResponse,
                    parse_obj_as(
                        type_=QuerySpansResponse,  # type: ignore
                        object_=_response.json(),
                    ),
                )
                _items = _parsed_response.data
                _has_next = False
                _get_next = None
                if _parsed_response.pagination is not None:
                    _parsed_next = _parsed_response.pagination.next_page_token
                    _has_next = _parsed_next is not None and _parsed_next != ""

                    async def _get_next():
                        return await self.query_spans(
                            start_time=start_time,
                            tracing_project_fqn=tracing_project_fqn,
                            end_time=end_time,
                            trace_ids=trace_ids,
                            span_ids=span_ids,
                            parent_span_ids=parent_span_ids,
                            created_by_subject_types=created_by_subject_types,
                            created_by_subject_slugs=created_by_subject_slugs,
                            application_names=application_names,
                            limit=limit,
                            sort_direction=sort_direction,
                            page_token=_parsed_next,
                            filters=filters,
                            request_options=request_options,
                        )

                return AsyncPager(has_next=_has_next, items=_items, get_next=_get_next, response=_parsed_response)
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
        raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
