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

import typing

from ..core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
from ..core.request_options import RequestOptions
from ..types.focus import Focus
from ..types.format import Format
from ..types.o_tel_flat_span import OTelFlatSpan
from ..types.o_tel_links_response import OTelLinksResponse
from ..types.o_tel_spans_tree import OTelSpansTree
from ..types.o_tel_tracing_response import OTelTracingResponse
from .raw_client import AsyncRawTracingClient, RawTracingClient
from .types.fetch_trace_request_trace_id import FetchTraceRequestTraceId
from .types.remove_trace_request_trace_id import RemoveTraceRequestTraceId

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


class TracingClient:
    def __init__(self, *, client_wrapper: SyncClientWrapper):
        self._raw_client = RawTracingClient(client_wrapper=client_wrapper)

    @property
    def with_raw_response(self) -> RawTracingClient:
        """
        Retrieves a raw implementation of this client that returns raw responses.

        Returns
        -------
        RawTracingClient
        """
        return self._raw_client

    def add_trace(
        self,
        *,
        spans: typing.Optional[typing.Sequence[OTelFlatSpan]] = OMIT,
        traces: typing.Optional[
            typing.Dict[str, typing.Optional[OTelSpansTree]]
        ] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> OTelLinksResponse:
        """
        Parameters
        ----------
        spans : typing.Optional[typing.Sequence[OTelFlatSpan]]

        traces : typing.Optional[typing.Dict[str, typing.Optional[OTelSpansTree]]]

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

        Returns
        -------
        OTelLinksResponse
            Successful Response

        Examples
        --------
        from agenta import AgentaApi

        client = AgentaApi(
            api_key="YOUR_API_KEY",
        )
        client.tracing.add_trace()
        """
        _response = self._raw_client.add_trace(
            spans=spans, traces=traces, request_options=request_options
        )
        return _response.data

    def edit_trace(
        self,
        *,
        spans: typing.Optional[typing.Sequence[OTelFlatSpan]] = OMIT,
        traces: typing.Optional[
            typing.Dict[str, typing.Optional[OTelSpansTree]]
        ] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> OTelLinksResponse:
        """
        Parameters
        ----------
        spans : typing.Optional[typing.Sequence[OTelFlatSpan]]

        traces : typing.Optional[typing.Dict[str, typing.Optional[OTelSpansTree]]]

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

        Returns
        -------
        OTelLinksResponse
            Successful Response

        Examples
        --------
        from agenta import AgentaApi

        client = AgentaApi(
            api_key="YOUR_API_KEY",
        )
        client.tracing.edit_trace()
        """
        _response = self._raw_client.edit_trace(
            spans=spans, traces=traces, request_options=request_options
        )
        return _response.data

    def fetch_trace(
        self,
        trace_id: FetchTraceRequestTraceId,
        *,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> OTelTracingResponse:
        """
        Parameters
        ----------
        trace_id : FetchTraceRequestTraceId

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

        Returns
        -------
        OTelTracingResponse
            Successful Response

        Examples
        --------
        from agenta import AgentaApi

        client = AgentaApi(
            api_key="YOUR_API_KEY",
        )
        client.tracing.fetch_trace(
            trace_id="trace_id",
        )
        """
        _response = self._raw_client.fetch_trace(
            trace_id, request_options=request_options
        )
        return _response.data

    def remove_trace(
        self,
        trace_id: RemoveTraceRequestTraceId,
        *,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> OTelLinksResponse:
        """
        Parameters
        ----------
        trace_id : RemoveTraceRequestTraceId

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

        Returns
        -------
        OTelLinksResponse
            Successful Response

        Examples
        --------
        from agenta import AgentaApi

        client = AgentaApi(
            api_key="YOUR_API_KEY",
        )
        client.tracing.remove_trace(
            trace_id="trace_id",
        )
        """
        _response = self._raw_client.remove_trace(
            trace_id, request_options=request_options
        )
        return _response.data

    def query_spans(
        self,
        *,
        filter: str,
        focus: typing.Optional[Focus] = None,
        format: typing.Optional[Format] = None,
        oldest: typing.Optional[str] = None,
        newest: typing.Optional[str] = None,
        limit: typing.Optional[int] = None,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> OTelTracingResponse:
        """
        Parameters
        ----------
        filter : str

        focus : typing.Optional[Focus]

        format : typing.Optional[Format]

        oldest : typing.Optional[str]

        newest : typing.Optional[str]

        limit : typing.Optional[int]

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

        Returns
        -------
        OTelTracingResponse
            Successful Response

        Examples
        --------
        from agenta import AgentaApi

        client = AgentaApi(
            api_key="YOUR_API_KEY",
        )
        client.tracing.query_spans(
            filter="filter",
        )
        """
        _response = self._raw_client.query_spans(
            filter=filter,
            focus=focus,
            format=format,
            oldest=oldest,
            newest=newest,
            limit=limit,
            request_options=request_options,
        )
        return _response.data

    def ingest_spans(
        self,
        *,
        spans: typing.Optional[typing.Sequence[OTelFlatSpan]] = OMIT,
        traces: typing.Optional[
            typing.Dict[str, typing.Optional[OTelSpansTree]]
        ] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> OTelLinksResponse:
        """
        Parameters
        ----------
        spans : typing.Optional[typing.Sequence[OTelFlatSpan]]

        traces : typing.Optional[typing.Dict[str, typing.Optional[OTelSpansTree]]]

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

        Returns
        -------
        OTelLinksResponse
            Successful Response

        Examples
        --------
        from agenta import AgentaApi

        client = AgentaApi(
            api_key="YOUR_API_KEY",
        )
        client.tracing.ingest_spans()
        """
        _response = self._raw_client.ingest_spans(
            spans=spans, traces=traces, request_options=request_options
        )
        return _response.data

    def ingest_spans_rpc(
        self,
        *,
        spans: typing.Optional[typing.Sequence[OTelFlatSpan]] = OMIT,
        traces: typing.Optional[
            typing.Dict[str, typing.Optional[OTelSpansTree]]
        ] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> OTelLinksResponse:
        """
        Parameters
        ----------
        spans : typing.Optional[typing.Sequence[OTelFlatSpan]]

        traces : typing.Optional[typing.Dict[str, typing.Optional[OTelSpansTree]]]

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

        Returns
        -------
        OTelLinksResponse
            Successful Response

        Examples
        --------
        from agenta import AgentaApi

        client = AgentaApi(
            api_key="YOUR_API_KEY",
        )
        client.tracing.ingest_spans_rpc()
        """
        _response = self._raw_client.ingest_spans_rpc(
            spans=spans, traces=traces, request_options=request_options
        )
        return _response.data

    def query_spans_rpc(
        self,
        *,
        filter: str,
        focus: typing.Optional[Focus] = None,
        format: typing.Optional[Format] = None,
        oldest: typing.Optional[str] = None,
        newest: typing.Optional[str] = None,
        limit: typing.Optional[int] = None,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> OTelTracingResponse:
        """
        Parameters
        ----------
        filter : str

        focus : typing.Optional[Focus]

        format : typing.Optional[Format]

        oldest : typing.Optional[str]

        newest : typing.Optional[str]

        limit : typing.Optional[int]

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

        Returns
        -------
        OTelTracingResponse
            Successful Response

        Examples
        --------
        from agenta import AgentaApi

        client = AgentaApi(
            api_key="YOUR_API_KEY",
        )
        client.tracing.query_spans_rpc(
            filter="filter",
        )
        """
        _response = self._raw_client.query_spans_rpc(
            filter=filter,
            focus=focus,
            format=format,
            oldest=oldest,
            newest=newest,
            limit=limit,
            request_options=request_options,
        )
        return _response.data


class AsyncTracingClient:
    def __init__(self, *, client_wrapper: AsyncClientWrapper):
        self._raw_client = AsyncRawTracingClient(client_wrapper=client_wrapper)

    @property
    def with_raw_response(self) -> AsyncRawTracingClient:
        """
        Retrieves a raw implementation of this client that returns raw responses.

        Returns
        -------
        AsyncRawTracingClient
        """
        return self._raw_client

    async def add_trace(
        self,
        *,
        spans: typing.Optional[typing.Sequence[OTelFlatSpan]] = OMIT,
        traces: typing.Optional[
            typing.Dict[str, typing.Optional[OTelSpansTree]]
        ] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> OTelLinksResponse:
        """
        Parameters
        ----------
        spans : typing.Optional[typing.Sequence[OTelFlatSpan]]

        traces : typing.Optional[typing.Dict[str, typing.Optional[OTelSpansTree]]]

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

        Returns
        -------
        OTelLinksResponse
            Successful Response

        Examples
        --------
        import asyncio

        from agenta import AsyncAgentaApi

        client = AsyncAgentaApi(
            api_key="YOUR_API_KEY",
        )


        async def main() -> None:
            await client.tracing.add_trace()


        asyncio.run(main())
        """
        _response = await self._raw_client.add_trace(
            spans=spans, traces=traces, request_options=request_options
        )
        return _response.data

    async def edit_trace(
        self,
        *,
        spans: typing.Optional[typing.Sequence[OTelFlatSpan]] = OMIT,
        traces: typing.Optional[
            typing.Dict[str, typing.Optional[OTelSpansTree]]
        ] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> OTelLinksResponse:
        """
        Parameters
        ----------
        spans : typing.Optional[typing.Sequence[OTelFlatSpan]]

        traces : typing.Optional[typing.Dict[str, typing.Optional[OTelSpansTree]]]

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

        Returns
        -------
        OTelLinksResponse
            Successful Response

        Examples
        --------
        import asyncio

        from agenta import AsyncAgentaApi

        client = AsyncAgentaApi(
            api_key="YOUR_API_KEY",
        )


        async def main() -> None:
            await client.tracing.edit_trace()


        asyncio.run(main())
        """
        _response = await self._raw_client.edit_trace(
            spans=spans, traces=traces, request_options=request_options
        )
        return _response.data

    async def fetch_trace(
        self,
        trace_id: FetchTraceRequestTraceId,
        *,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> OTelTracingResponse:
        """
        Parameters
        ----------
        trace_id : FetchTraceRequestTraceId

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

        Returns
        -------
        OTelTracingResponse
            Successful Response

        Examples
        --------
        import asyncio

        from agenta import AsyncAgentaApi

        client = AsyncAgentaApi(
            api_key="YOUR_API_KEY",
        )


        async def main() -> None:
            await client.tracing.fetch_trace(
                trace_id="trace_id",
            )


        asyncio.run(main())
        """
        _response = await self._raw_client.fetch_trace(
            trace_id, request_options=request_options
        )
        return _response.data

    async def remove_trace(
        self,
        trace_id: RemoveTraceRequestTraceId,
        *,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> OTelLinksResponse:
        """
        Parameters
        ----------
        trace_id : RemoveTraceRequestTraceId

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

        Returns
        -------
        OTelLinksResponse
            Successful Response

        Examples
        --------
        import asyncio

        from agenta import AsyncAgentaApi

        client = AsyncAgentaApi(
            api_key="YOUR_API_KEY",
        )


        async def main() -> None:
            await client.tracing.remove_trace(
                trace_id="trace_id",
            )


        asyncio.run(main())
        """
        _response = await self._raw_client.remove_trace(
            trace_id, request_options=request_options
        )
        return _response.data

    async def query_spans(
        self,
        *,
        filter: str,
        focus: typing.Optional[Focus] = None,
        format: typing.Optional[Format] = None,
        oldest: typing.Optional[str] = None,
        newest: typing.Optional[str] = None,
        limit: typing.Optional[int] = None,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> OTelTracingResponse:
        """
        Parameters
        ----------
        filter : str

        focus : typing.Optional[Focus]

        format : typing.Optional[Format]

        oldest : typing.Optional[str]

        newest : typing.Optional[str]

        limit : typing.Optional[int]

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

        Returns
        -------
        OTelTracingResponse
            Successful Response

        Examples
        --------
        import asyncio

        from agenta import AsyncAgentaApi

        client = AsyncAgentaApi(
            api_key="YOUR_API_KEY",
        )


        async def main() -> None:
            await client.tracing.query_spans(
                filter="filter",
            )


        asyncio.run(main())
        """
        _response = await self._raw_client.query_spans(
            filter=filter,
            focus=focus,
            format=format,
            oldest=oldest,
            newest=newest,
            limit=limit,
            request_options=request_options,
        )
        return _response.data

    async def ingest_spans(
        self,
        *,
        spans: typing.Optional[typing.Sequence[OTelFlatSpan]] = OMIT,
        traces: typing.Optional[
            typing.Dict[str, typing.Optional[OTelSpansTree]]
        ] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> OTelLinksResponse:
        """
        Parameters
        ----------
        spans : typing.Optional[typing.Sequence[OTelFlatSpan]]

        traces : typing.Optional[typing.Dict[str, typing.Optional[OTelSpansTree]]]

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

        Returns
        -------
        OTelLinksResponse
            Successful Response

        Examples
        --------
        import asyncio

        from agenta import AsyncAgentaApi

        client = AsyncAgentaApi(
            api_key="YOUR_API_KEY",
        )


        async def main() -> None:
            await client.tracing.ingest_spans()


        asyncio.run(main())
        """
        _response = await self._raw_client.ingest_spans(
            spans=spans, traces=traces, request_options=request_options
        )
        return _response.data

    async def ingest_spans_rpc(
        self,
        *,
        spans: typing.Optional[typing.Sequence[OTelFlatSpan]] = OMIT,
        traces: typing.Optional[
            typing.Dict[str, typing.Optional[OTelSpansTree]]
        ] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> OTelLinksResponse:
        """
        Parameters
        ----------
        spans : typing.Optional[typing.Sequence[OTelFlatSpan]]

        traces : typing.Optional[typing.Dict[str, typing.Optional[OTelSpansTree]]]

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

        Returns
        -------
        OTelLinksResponse
            Successful Response

        Examples
        --------
        import asyncio

        from agenta import AsyncAgentaApi

        client = AsyncAgentaApi(
            api_key="YOUR_API_KEY",
        )


        async def main() -> None:
            await client.tracing.ingest_spans_rpc()


        asyncio.run(main())
        """
        _response = await self._raw_client.ingest_spans_rpc(
            spans=spans, traces=traces, request_options=request_options
        )
        return _response.data

    async def query_spans_rpc(
        self,
        *,
        filter: str,
        focus: typing.Optional[Focus] = None,
        format: typing.Optional[Format] = None,
        oldest: typing.Optional[str] = None,
        newest: typing.Optional[str] = None,
        limit: typing.Optional[int] = None,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> OTelTracingResponse:
        """
        Parameters
        ----------
        filter : str

        focus : typing.Optional[Focus]

        format : typing.Optional[Format]

        oldest : typing.Optional[str]

        newest : typing.Optional[str]

        limit : typing.Optional[int]

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

        Returns
        -------
        OTelTracingResponse
            Successful Response

        Examples
        --------
        import asyncio

        from agenta import AsyncAgentaApi

        client = AsyncAgentaApi(
            api_key="YOUR_API_KEY",
        )


        async def main() -> None:
            await client.tracing.query_spans_rpc(
                filter="filter",
            )


        asyncio.run(main())
        """
        _response = await self._raw_client.query_spans_rpc(
            filter=filter,
            focus=focus,
            format=format,
            oldest=oldest,
            newest=newest,
            limit=limit,
            request_options=request_options,
        )
        return _response.data
