# 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.big_query_dataset import BigQueryDataset
from ..types.big_query_table import BigQueryTable
from ..types.compute_instance import ComputeInstance
from ..types.create_sandbox_response import CreateSandboxResponse
from ..types.dump_sandbox_response_google_cloud_data import DumpSandboxResponseGoogleCloudData
from ..types.initialize_sandbox_response import InitializeSandboxResponse
from ..types.log_bucket import LogBucket
from ..types.log_entry import LogEntry
from ..types.log_sink import LogSink
from ..types.release_sandbox_response import ReleaseSandboxResponse
from ..types.reset_sandbox_response import ResetSandboxResponse
from ..types.sandbox_info import SandboxInfo
from ..types.sandbox_mcp_server import SandboxMcpServer
from ..types.storage_bucket import StorageBucket
from ..types.storage_object import StorageObject
from .raw_client import AsyncRawSandboxClient, RawSandboxClient

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


class SandboxClient:
    def __init__(self, *, client_wrapper: SyncClientWrapper):
        self._raw_client = RawSandboxClient(client_wrapper=client_wrapper)

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

        Returns
        -------
        RawSandboxClient
        """
        return self._raw_client

    def create_sandbox(
        self, server_name: SandboxMcpServer, *, request_options: typing.Optional[RequestOptions] = None
    ) -> CreateSandboxResponse:
        """
        Acquire an idle sandbox instance for a specific MCP server. The sandbox will be marked as 'occupied'.

        Parameters
        ----------
        server_name : SandboxMcpServer
            The MCP server name

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

        Returns
        -------
        CreateSandboxResponse
            Successful Response

        Examples
        --------
        from klavis import Klavis, SandboxMcpServer

        client = Klavis(
            api_key="YOUR_API_KEY",
        )
        client.sandbox.create_sandbox(
            server_name=SandboxMcpServer.JIRA,
        )
        """
        _response = self._raw_client.create_sandbox(server_name, request_options=request_options)
        return _response.data

    def get_sandbox(
        self, server_name: SandboxMcpServer, sandbox_id: str, *, request_options: typing.Optional[RequestOptions] = None
    ) -> SandboxInfo:
        """
        Retrieve detailed information about a specific sandbox instance.

        Parameters
        ----------
        server_name : SandboxMcpServer
            The MCP server name

        sandbox_id : str
            The unique sandbox identifier

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

        Returns
        -------
        SandboxInfo
            Successful Response

        Examples
        --------
        from klavis import Klavis, SandboxMcpServer

        client = Klavis(
            api_key="YOUR_API_KEY",
        )
        client.sandbox.get_sandbox(
            server_name=SandboxMcpServer.JIRA,
            sandbox_id="sandbox_id",
        )
        """
        _response = self._raw_client.get_sandbox(server_name, sandbox_id, request_options=request_options)
        return _response.data

    def delete_sandbox(
        self, server_name: SandboxMcpServer, sandbox_id: str, *, request_options: typing.Optional[RequestOptions] = None
    ) -> ReleaseSandboxResponse:
        """
        Release an occupied sandbox back to idle state and marks the sandbox as available for reuse.

        Parameters
        ----------
        server_name : SandboxMcpServer
            The MCP server name

        sandbox_id : str
            The unique sandbox identifier

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

        Returns
        -------
        ReleaseSandboxResponse
            Successful Response

        Examples
        --------
        from klavis import Klavis, SandboxMcpServer

        client = Klavis(
            api_key="YOUR_API_KEY",
        )
        client.sandbox.delete_sandbox(
            server_name=SandboxMcpServer.JIRA,
            sandbox_id="sandbox_id",
        )
        """
        _response = self._raw_client.delete_sandbox(server_name, sandbox_id, request_options=request_options)
        return _response.data

    def reset_sandbox(
        self, server_name: SandboxMcpServer, sandbox_id: str, *, request_options: typing.Optional[RequestOptions] = None
    ) -> ResetSandboxResponse:
        """
        Reset the sandbox to its initial empty state, clearing all data while maintaining the sandbox instance.

        Parameters
        ----------
        server_name : SandboxMcpServer
            The MCP server name

        sandbox_id : str
            The unique sandbox identifier

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

        Returns
        -------
        ResetSandboxResponse
            Successful Response

        Examples
        --------
        from klavis import Klavis, SandboxMcpServer

        client = Klavis(
            api_key="YOUR_API_KEY",
        )
        client.sandbox.reset_sandbox(
            server_name=SandboxMcpServer.JIRA,
            sandbox_id="sandbox_id",
        )
        """
        _response = self._raw_client.reset_sandbox(server_name, sandbox_id, request_options=request_options)
        return _response.data

    def initialize_sandbox(
        self,
        sandbox_id: str,
        *,
        datasets: typing.Optional[typing.Sequence[BigQueryDataset]] = OMIT,
        tables: typing.Optional[typing.Sequence[BigQueryTable]] = OMIT,
        buckets: typing.Optional[typing.Sequence[StorageBucket]] = OMIT,
        objects: typing.Optional[typing.Sequence[StorageObject]] = OMIT,
        log_entries: typing.Optional[typing.Sequence[LogEntry]] = OMIT,
        log_sinks: typing.Optional[typing.Sequence[LogSink]] = OMIT,
        log_buckets: typing.Optional[typing.Sequence[LogBucket]] = OMIT,
        instances: typing.Optional[typing.Sequence[ComputeInstance]] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> InitializeSandboxResponse:
        """
        Initialize the sandbox with google_cloud-specific data following the defined schema.

        Parameters
        ----------
        sandbox_id : str
            The unique sandbox identifier

        datasets : typing.Optional[typing.Sequence[BigQueryDataset]]
            BigQuery datasets

        tables : typing.Optional[typing.Sequence[BigQueryTable]]
            BigQuery tables

        buckets : typing.Optional[typing.Sequence[StorageBucket]]
            Cloud Storage buckets

        objects : typing.Optional[typing.Sequence[StorageObject]]
            Cloud Storage objects

        log_entries : typing.Optional[typing.Sequence[LogEntry]]
            Log entries

        log_sinks : typing.Optional[typing.Sequence[LogSink]]
            Log sinks

        log_buckets : typing.Optional[typing.Sequence[LogBucket]]
            Log buckets

        instances : typing.Optional[typing.Sequence[ComputeInstance]]
            Compute Engine instances

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

        Returns
        -------
        InitializeSandboxResponse
            Successful Response

        Examples
        --------
        from klavis import Klavis

        client = Klavis(
            api_key="YOUR_API_KEY",
        )
        client.sandbox.initialize_sandbox(
            sandbox_id="sandbox_id",
        )
        """
        _response = self._raw_client.initialize_sandbox(
            sandbox_id,
            datasets=datasets,
            tables=tables,
            buckets=buckets,
            objects=objects,
            log_entries=log_entries,
            log_sinks=log_sinks,
            log_buckets=log_buckets,
            instances=instances,
            request_options=request_options,
        )
        return _response.data

    def dump_sandbox(
        self, sandbox_id: str, *, request_options: typing.Optional[RequestOptions] = None
    ) -> DumpSandboxResponseGoogleCloudData:
        """
        Export all data from the sandbox in the same format used for initialization.

        Parameters
        ----------
        sandbox_id : str
            The unique sandbox identifier

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

        Returns
        -------
        DumpSandboxResponseGoogleCloudData
            Successful Response

        Examples
        --------
        from klavis import Klavis

        client = Klavis(
            api_key="YOUR_API_KEY",
        )
        client.sandbox.dump_sandbox(
            sandbox_id="sandbox_id",
        )
        """
        _response = self._raw_client.dump_sandbox(sandbox_id, request_options=request_options)
        return _response.data


class AsyncSandboxClient:
    def __init__(self, *, client_wrapper: AsyncClientWrapper):
        self._raw_client = AsyncRawSandboxClient(client_wrapper=client_wrapper)

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

        Returns
        -------
        AsyncRawSandboxClient
        """
        return self._raw_client

    async def create_sandbox(
        self, server_name: SandboxMcpServer, *, request_options: typing.Optional[RequestOptions] = None
    ) -> CreateSandboxResponse:
        """
        Acquire an idle sandbox instance for a specific MCP server. The sandbox will be marked as 'occupied'.

        Parameters
        ----------
        server_name : SandboxMcpServer
            The MCP server name

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

        Returns
        -------
        CreateSandboxResponse
            Successful Response

        Examples
        --------
        import asyncio

        from klavis import AsyncKlavis, SandboxMcpServer

        client = AsyncKlavis(
            api_key="YOUR_API_KEY",
        )


        async def main() -> None:
            await client.sandbox.create_sandbox(
                server_name=SandboxMcpServer.JIRA,
            )


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

    async def get_sandbox(
        self, server_name: SandboxMcpServer, sandbox_id: str, *, request_options: typing.Optional[RequestOptions] = None
    ) -> SandboxInfo:
        """
        Retrieve detailed information about a specific sandbox instance.

        Parameters
        ----------
        server_name : SandboxMcpServer
            The MCP server name

        sandbox_id : str
            The unique sandbox identifier

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

        Returns
        -------
        SandboxInfo
            Successful Response

        Examples
        --------
        import asyncio

        from klavis import AsyncKlavis, SandboxMcpServer

        client = AsyncKlavis(
            api_key="YOUR_API_KEY",
        )


        async def main() -> None:
            await client.sandbox.get_sandbox(
                server_name=SandboxMcpServer.JIRA,
                sandbox_id="sandbox_id",
            )


        asyncio.run(main())
        """
        _response = await self._raw_client.get_sandbox(server_name, sandbox_id, request_options=request_options)
        return _response.data

    async def delete_sandbox(
        self, server_name: SandboxMcpServer, sandbox_id: str, *, request_options: typing.Optional[RequestOptions] = None
    ) -> ReleaseSandboxResponse:
        """
        Release an occupied sandbox back to idle state and marks the sandbox as available for reuse.

        Parameters
        ----------
        server_name : SandboxMcpServer
            The MCP server name

        sandbox_id : str
            The unique sandbox identifier

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

        Returns
        -------
        ReleaseSandboxResponse
            Successful Response

        Examples
        --------
        import asyncio

        from klavis import AsyncKlavis, SandboxMcpServer

        client = AsyncKlavis(
            api_key="YOUR_API_KEY",
        )


        async def main() -> None:
            await client.sandbox.delete_sandbox(
                server_name=SandboxMcpServer.JIRA,
                sandbox_id="sandbox_id",
            )


        asyncio.run(main())
        """
        _response = await self._raw_client.delete_sandbox(server_name, sandbox_id, request_options=request_options)
        return _response.data

    async def reset_sandbox(
        self, server_name: SandboxMcpServer, sandbox_id: str, *, request_options: typing.Optional[RequestOptions] = None
    ) -> ResetSandboxResponse:
        """
        Reset the sandbox to its initial empty state, clearing all data while maintaining the sandbox instance.

        Parameters
        ----------
        server_name : SandboxMcpServer
            The MCP server name

        sandbox_id : str
            The unique sandbox identifier

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

        Returns
        -------
        ResetSandboxResponse
            Successful Response

        Examples
        --------
        import asyncio

        from klavis import AsyncKlavis, SandboxMcpServer

        client = AsyncKlavis(
            api_key="YOUR_API_KEY",
        )


        async def main() -> None:
            await client.sandbox.reset_sandbox(
                server_name=SandboxMcpServer.JIRA,
                sandbox_id="sandbox_id",
            )


        asyncio.run(main())
        """
        _response = await self._raw_client.reset_sandbox(server_name, sandbox_id, request_options=request_options)
        return _response.data

    async def initialize_sandbox(
        self,
        sandbox_id: str,
        *,
        datasets: typing.Optional[typing.Sequence[BigQueryDataset]] = OMIT,
        tables: typing.Optional[typing.Sequence[BigQueryTable]] = OMIT,
        buckets: typing.Optional[typing.Sequence[StorageBucket]] = OMIT,
        objects: typing.Optional[typing.Sequence[StorageObject]] = OMIT,
        log_entries: typing.Optional[typing.Sequence[LogEntry]] = OMIT,
        log_sinks: typing.Optional[typing.Sequence[LogSink]] = OMIT,
        log_buckets: typing.Optional[typing.Sequence[LogBucket]] = OMIT,
        instances: typing.Optional[typing.Sequence[ComputeInstance]] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> InitializeSandboxResponse:
        """
        Initialize the sandbox with google_cloud-specific data following the defined schema.

        Parameters
        ----------
        sandbox_id : str
            The unique sandbox identifier

        datasets : typing.Optional[typing.Sequence[BigQueryDataset]]
            BigQuery datasets

        tables : typing.Optional[typing.Sequence[BigQueryTable]]
            BigQuery tables

        buckets : typing.Optional[typing.Sequence[StorageBucket]]
            Cloud Storage buckets

        objects : typing.Optional[typing.Sequence[StorageObject]]
            Cloud Storage objects

        log_entries : typing.Optional[typing.Sequence[LogEntry]]
            Log entries

        log_sinks : typing.Optional[typing.Sequence[LogSink]]
            Log sinks

        log_buckets : typing.Optional[typing.Sequence[LogBucket]]
            Log buckets

        instances : typing.Optional[typing.Sequence[ComputeInstance]]
            Compute Engine instances

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

        Returns
        -------
        InitializeSandboxResponse
            Successful Response

        Examples
        --------
        import asyncio

        from klavis import AsyncKlavis

        client = AsyncKlavis(
            api_key="YOUR_API_KEY",
        )


        async def main() -> None:
            await client.sandbox.initialize_sandbox(
                sandbox_id="sandbox_id",
            )


        asyncio.run(main())
        """
        _response = await self._raw_client.initialize_sandbox(
            sandbox_id,
            datasets=datasets,
            tables=tables,
            buckets=buckets,
            objects=objects,
            log_entries=log_entries,
            log_sinks=log_sinks,
            log_buckets=log_buckets,
            instances=instances,
            request_options=request_options,
        )
        return _response.data

    async def dump_sandbox(
        self, sandbox_id: str, *, request_options: typing.Optional[RequestOptions] = None
    ) -> DumpSandboxResponseGoogleCloudData:
        """
        Export all data from the sandbox in the same format used for initialization.

        Parameters
        ----------
        sandbox_id : str
            The unique sandbox identifier

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

        Returns
        -------
        DumpSandboxResponseGoogleCloudData
            Successful Response

        Examples
        --------
        import asyncio

        from klavis import AsyncKlavis

        client = AsyncKlavis(
            api_key="YOUR_API_KEY",
        )


        async def main() -> None:
            await client.sandbox.dump_sandbox(
                sandbox_id="sandbox_id",
            )


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