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

import typing
import urllib.parse
from json.decoder import JSONDecodeError

from ...core.api_error import ApiError
from ...core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
from ...core.jsonable_encoder import jsonable_encoder
from ...core.remove_none_from_dict import remove_none_from_dict
from ...errors.unprocessable_entity_error import UnprocessableEntityError
from ...types.batch import Batch
from ...types.batch_paginated_list import BatchPaginatedList
from ...types.batch_public_output import BatchPublicOutput
from ...types.http_validation_error import HttpValidationError
from ...types.llama_parse_parameters import LlamaParseParameters

try:
    import pydantic
    if pydantic.__version__.startswith("1."):
        raise ImportError
    import pydantic.v1 as pydantic  # type: ignore
except ImportError:
    import pydantic  # type: ignore

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


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

    def list_batches(
        self,
        *,
        limit: typing.Optional[int] = None,
        offset: typing.Optional[int] = None,
        project_id: typing.Optional[str] = None,
        organization_id: typing.Optional[str] = None,
    ) -> BatchPaginatedList:
        """
        Parameters:
            - limit: typing.Optional[int].

            - offset: typing.Optional[int].

            - project_id: typing.Optional[str].

            - organization_id: typing.Optional[str].
        ---
        from llama_cloud.client import LlamaCloud

        client = LlamaCloud(
            token="YOUR_TOKEN",
        )
        client.beta.list_batches()
        """
        _response = self._client_wrapper.httpx_client.request(
            "GET",
            urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v1/beta/batches"),
            params=remove_none_from_dict(
                {"limit": limit, "offset": offset, "project_id": project_id, "organization_id": organization_id}
            ),
            headers=self._client_wrapper.get_headers(),
            timeout=60,
        )
        if 200 <= _response.status_code < 300:
            return pydantic.parse_obj_as(BatchPaginatedList, _response.json())  # type: ignore
        if _response.status_code == 422:
            raise UnprocessableEntityError(pydantic.parse_obj_as(HttpValidationError, _response.json()))  # type: ignore
        try:
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, body=_response.text)
        raise ApiError(status_code=_response.status_code, body=_response_json)

    def create_batch(
        self,
        *,
        organization_id: typing.Optional[str] = None,
        project_id: typing.Optional[str] = None,
        tool: str,
        tool_data: typing.Optional[LlamaParseParameters] = OMIT,
        input_type: str,
        input_id: str,
        output_type: typing.Optional[str] = OMIT,
        output_id: typing.Optional[str] = OMIT,
        batch_create_project_id: str,
        external_id: str,
        completion_window: typing.Optional[int] = OMIT,
    ) -> Batch:
        """
        Parameters:
            - organization_id: typing.Optional[str].

            - project_id: typing.Optional[str].

            - tool: str. The tool to be used for all requests in the batch.

            - tool_data: typing.Optional[LlamaParseParameters].

            - input_type: str. The type of input file. Currently only 'datasource' is supported.

            - input_id: str. The ID of the input file for the batch.

            - output_type: typing.Optional[str].

            - output_id: typing.Optional[str].

            - batch_create_project_id: str. The ID of the project to which the batch belongs

            - external_id: str. A developer-provided ID for the batch. This ID will be returned in the response.

            - completion_window: typing.Optional[int]. The time frame within which the batch should be processed. Currently only 24h is supported.
        ---
        from llama_cloud import FailPageMode, LlamaParseParameters, ParsingMode
        from llama_cloud.client import LlamaCloud

        client = LlamaCloud(
            token="YOUR_TOKEN",
        )
        client.beta.create_batch(
            tool="string",
            tool_data=LlamaParseParameters(
                parse_mode=ParsingMode.PARSE_PAGE_WITHOUT_LLM,
                replace_failed_page_mode=FailPageMode.RAW_TEXT,
            ),
            input_type="string",
            input_id="string",
            batch_create_project_id="string",
            external_id="string",
        )
        """
        _request: typing.Dict[str, typing.Any] = {
            "tool": tool,
            "input_type": input_type,
            "input_id": input_id,
            "project_id": batch_create_project_id,
            "external_id": external_id,
        }
        if tool_data is not OMIT:
            _request["tool_data"] = tool_data
        if output_type is not OMIT:
            _request["output_type"] = output_type
        if output_id is not OMIT:
            _request["output_id"] = output_id
        if completion_window is not OMIT:
            _request["completion_window"] = completion_window
        _response = self._client_wrapper.httpx_client.request(
            "POST",
            urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v1/beta/batches"),
            params=remove_none_from_dict({"organization_id": organization_id, "project_id": project_id}),
            json=jsonable_encoder(_request),
            headers=self._client_wrapper.get_headers(),
            timeout=60,
        )
        if 200 <= _response.status_code < 300:
            return pydantic.parse_obj_as(Batch, _response.json())  # type: ignore
        if _response.status_code == 422:
            raise UnprocessableEntityError(pydantic.parse_obj_as(HttpValidationError, _response.json()))  # type: ignore
        try:
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, body=_response.text)
        raise ApiError(status_code=_response.status_code, body=_response_json)

    def get_batch(self, batch_id: str, *, organization_id: typing.Optional[str] = None) -> BatchPublicOutput:
        """
        Parameters:
            - batch_id: str.

            - organization_id: typing.Optional[str].
        ---
        from llama_cloud.client import LlamaCloud

        client = LlamaCloud(
            token="YOUR_TOKEN",
        )
        client.beta.get_batch(
            batch_id="string",
        )
        """
        _response = self._client_wrapper.httpx_client.request(
            "GET",
            urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", f"api/v1/beta/batches/{batch_id}"),
            params=remove_none_from_dict({"organization_id": organization_id}),
            headers=self._client_wrapper.get_headers(),
            timeout=60,
        )
        if 200 <= _response.status_code < 300:
            return pydantic.parse_obj_as(BatchPublicOutput, _response.json())  # type: ignore
        if _response.status_code == 422:
            raise UnprocessableEntityError(pydantic.parse_obj_as(HttpValidationError, _response.json()))  # type: ignore
        try:
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, body=_response.text)
        raise ApiError(status_code=_response.status_code, body=_response_json)


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

    async def list_batches(
        self,
        *,
        limit: typing.Optional[int] = None,
        offset: typing.Optional[int] = None,
        project_id: typing.Optional[str] = None,
        organization_id: typing.Optional[str] = None,
    ) -> BatchPaginatedList:
        """
        Parameters:
            - limit: typing.Optional[int].

            - offset: typing.Optional[int].

            - project_id: typing.Optional[str].

            - organization_id: typing.Optional[str].
        ---
        from llama_cloud.client import AsyncLlamaCloud

        client = AsyncLlamaCloud(
            token="YOUR_TOKEN",
        )
        await client.beta.list_batches()
        """
        _response = await self._client_wrapper.httpx_client.request(
            "GET",
            urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v1/beta/batches"),
            params=remove_none_from_dict(
                {"limit": limit, "offset": offset, "project_id": project_id, "organization_id": organization_id}
            ),
            headers=self._client_wrapper.get_headers(),
            timeout=60,
        )
        if 200 <= _response.status_code < 300:
            return pydantic.parse_obj_as(BatchPaginatedList, _response.json())  # type: ignore
        if _response.status_code == 422:
            raise UnprocessableEntityError(pydantic.parse_obj_as(HttpValidationError, _response.json()))  # type: ignore
        try:
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, body=_response.text)
        raise ApiError(status_code=_response.status_code, body=_response_json)

    async def create_batch(
        self,
        *,
        organization_id: typing.Optional[str] = None,
        project_id: typing.Optional[str] = None,
        tool: str,
        tool_data: typing.Optional[LlamaParseParameters] = OMIT,
        input_type: str,
        input_id: str,
        output_type: typing.Optional[str] = OMIT,
        output_id: typing.Optional[str] = OMIT,
        batch_create_project_id: str,
        external_id: str,
        completion_window: typing.Optional[int] = OMIT,
    ) -> Batch:
        """
        Parameters:
            - organization_id: typing.Optional[str].

            - project_id: typing.Optional[str].

            - tool: str. The tool to be used for all requests in the batch.

            - tool_data: typing.Optional[LlamaParseParameters].

            - input_type: str. The type of input file. Currently only 'datasource' is supported.

            - input_id: str. The ID of the input file for the batch.

            - output_type: typing.Optional[str].

            - output_id: typing.Optional[str].

            - batch_create_project_id: str. The ID of the project to which the batch belongs

            - external_id: str. A developer-provided ID for the batch. This ID will be returned in the response.

            - completion_window: typing.Optional[int]. The time frame within which the batch should be processed. Currently only 24h is supported.
        ---
        from llama_cloud import FailPageMode, LlamaParseParameters, ParsingMode
        from llama_cloud.client import AsyncLlamaCloud

        client = AsyncLlamaCloud(
            token="YOUR_TOKEN",
        )
        await client.beta.create_batch(
            tool="string",
            tool_data=LlamaParseParameters(
                parse_mode=ParsingMode.PARSE_PAGE_WITHOUT_LLM,
                replace_failed_page_mode=FailPageMode.RAW_TEXT,
            ),
            input_type="string",
            input_id="string",
            batch_create_project_id="string",
            external_id="string",
        )
        """
        _request: typing.Dict[str, typing.Any] = {
            "tool": tool,
            "input_type": input_type,
            "input_id": input_id,
            "project_id": batch_create_project_id,
            "external_id": external_id,
        }
        if tool_data is not OMIT:
            _request["tool_data"] = tool_data
        if output_type is not OMIT:
            _request["output_type"] = output_type
        if output_id is not OMIT:
            _request["output_id"] = output_id
        if completion_window is not OMIT:
            _request["completion_window"] = completion_window
        _response = await self._client_wrapper.httpx_client.request(
            "POST",
            urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v1/beta/batches"),
            params=remove_none_from_dict({"organization_id": organization_id, "project_id": project_id}),
            json=jsonable_encoder(_request),
            headers=self._client_wrapper.get_headers(),
            timeout=60,
        )
        if 200 <= _response.status_code < 300:
            return pydantic.parse_obj_as(Batch, _response.json())  # type: ignore
        if _response.status_code == 422:
            raise UnprocessableEntityError(pydantic.parse_obj_as(HttpValidationError, _response.json()))  # type: ignore
        try:
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, body=_response.text)
        raise ApiError(status_code=_response.status_code, body=_response_json)

    async def get_batch(self, batch_id: str, *, organization_id: typing.Optional[str] = None) -> BatchPublicOutput:
        """
        Parameters:
            - batch_id: str.

            - organization_id: typing.Optional[str].
        ---
        from llama_cloud.client import AsyncLlamaCloud

        client = AsyncLlamaCloud(
            token="YOUR_TOKEN",
        )
        await client.beta.get_batch(
            batch_id="string",
        )
        """
        _response = await self._client_wrapper.httpx_client.request(
            "GET",
            urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", f"api/v1/beta/batches/{batch_id}"),
            params=remove_none_from_dict({"organization_id": organization_id}),
            headers=self._client_wrapper.get_headers(),
            timeout=60,
        )
        if 200 <= _response.status_code < 300:
            return pydantic.parse_obj_as(BatchPublicOutput, _response.json())  # type: ignore
        if _response.status_code == 422:
            raise UnprocessableEntityError(pydantic.parse_obj_as(HttpValidationError, _response.json()))  # type: ignore
        try:
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, body=_response.text)
        raise ApiError(status_code=_response.status_code, body=_response_json)
