# ----------------------------------------------------------------------------
#
#  Welcome to Baml! To use this generated code, please run the following:
#
#  $ pip install baml
#
# ----------------------------------------------------------------------------

# This file was generated by BAML: please do not edit it. Instead, edit the
# BAML files and re-generate this code using: baml-cli generate
# baml-cli is available with the baml package.

import typing

import baml_py
import typing_extensions

from . import stream_types, type_builder, types
from .globals import DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME as __runtime__
from .parser import LlmResponseParser, LlmStreamParser
from .runtime import BamlCallOptions, DoNotUseDirectlyCallManager


class BamlAsyncClient:
    __options: DoNotUseDirectlyCallManager
    __stream_client: "BamlStreamClient"
    __http_request: "BamlHttpRequestClient"
    __http_stream_request: "BamlHttpStreamRequestClient"
    __llm_response_parser: LlmResponseParser
    __llm_stream_parser: LlmStreamParser

    def __init__(self, options: DoNotUseDirectlyCallManager):
        self.__options = options
        self.__stream_client = BamlStreamClient(options)
        self.__http_request = BamlHttpRequestClient(options)
        self.__http_stream_request = BamlHttpStreamRequestClient(options)
        self.__llm_response_parser = LlmResponseParser(options)
        self.__llm_stream_parser = LlmStreamParser(options)

    def with_options(
        self,
        tb: typing.Optional[type_builder.TypeBuilder] = None,
        client_registry: typing.Optional[baml_py.baml_py.ClientRegistry] = None,
        client: typing.Optional[str] = None,
        collector: typing.Optional[
            typing.Union[baml_py.baml_py.Collector, typing.List[baml_py.baml_py.Collector]]
        ] = None,
        env: typing.Optional[typing.Dict[str, typing.Optional[str]]] = None,
        tags: typing.Optional[typing.Dict[str, str]] = None,
        on_tick: typing.Optional[typing.Callable[[str, baml_py.baml_py.FunctionLog], None]] = None,
    ) -> "BamlAsyncClient":
        options: BamlCallOptions = {}
        if tb is not None:
            options["tb"] = tb
        if client_registry is not None:
            options["client_registry"] = client_registry
        if client is not None:
            options["client"] = client
        if collector is not None:
            options["collector"] = collector
        if env is not None:
            options["env"] = env
        if tags is not None:
            options["tags"] = tags
        if on_tick is not None:
            options["on_tick"] = on_tick
        return BamlAsyncClient(self.__options.merge_options(options))

    @property
    def stream(self):
        return self.__stream_client

    @property
    def request(self):
        return self.__http_request

    @property
    def stream_request(self):
        return self.__http_stream_request

    @property
    def parse(self):
        return self.__llm_response_parser

    @property
    def parse_stream(self):
        return self.__llm_stream_parser

    async def SummarizeChunkAlone(
        self,
        language: str,
        symbol: str,
        line_start: int,
        line_end: int,
        content: str,
        parent_context: str,
        file_context: str,
        baml_options: BamlCallOptions = {},
    ) -> types.ChunkSummary:
        # Check if on_tick is provided
        if "on_tick" in baml_options:
            # Use streaming internally when on_tick is provided
            __stream__ = self.stream.SummarizeChunkAlone(
                language=language,
                symbol=symbol,
                line_start=line_start,
                line_end=line_end,
                content=content,
                parent_context=parent_context,
                file_context=file_context,
                baml_options=baml_options,
            )
            return await __stream__.get_final_response()
        else:
            # Original non-streaming code
            __result__ = await self.__options.merge_options(baml_options).call_function_async(
                function_name="SummarizeChunkAlone",
                args={
                    "language": language,
                    "symbol": symbol,
                    "line_start": line_start,
                    "line_end": line_end,
                    "content": content,
                    "parent_context": parent_context,
                    "file_context": file_context,
                },
            )
            return typing.cast(
                types.ChunkSummary,
                __result__.cast_to(types, types, stream_types, False, __runtime__),
            )

    async def SummarizeFileChunks(
        self,
        language: str,
        file_path: str,
        file_content: str,
        chunks: typing.List["types.ChunkInput"],
        baml_options: BamlCallOptions = {},
    ) -> types.FileSummaryResult:
        # Check if on_tick is provided
        if "on_tick" in baml_options:
            # Use streaming internally when on_tick is provided
            __stream__ = self.stream.SummarizeFileChunks(
                language=language,
                file_path=file_path,
                file_content=file_content,
                chunks=chunks,
                baml_options=baml_options,
            )
            return await __stream__.get_final_response()
        else:
            # Original non-streaming code
            __result__ = await self.__options.merge_options(baml_options).call_function_async(
                function_name="SummarizeFileChunks",
                args={
                    "language": language,
                    "file_path": file_path,
                    "file_content": file_content,
                    "chunks": chunks,
                },
            )
            return typing.cast(
                types.FileSummaryResult,
                __result__.cast_to(types, types, stream_types, False, __runtime__),
            )


class BamlStreamClient:
    __options: DoNotUseDirectlyCallManager

    def __init__(self, options: DoNotUseDirectlyCallManager):
        self.__options = options

    def SummarizeChunkAlone(
        self,
        language: str,
        symbol: str,
        line_start: int,
        line_end: int,
        content: str,
        parent_context: str,
        file_context: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlStream[stream_types.ChunkSummary, types.ChunkSummary]:
        __ctx__, __result__ = self.__options.merge_options(baml_options).create_async_stream(
            function_name="SummarizeChunkAlone",
            args={
                "language": language,
                "symbol": symbol,
                "line_start": line_start,
                "line_end": line_end,
                "content": content,
                "parent_context": parent_context,
                "file_context": file_context,
            },
        )
        return baml_py.BamlStream[stream_types.ChunkSummary, types.ChunkSummary](
            __result__,
            lambda x: typing.cast(
                stream_types.ChunkSummary, x.cast_to(types, types, stream_types, True, __runtime__)
            ),
            lambda x: typing.cast(
                types.ChunkSummary, x.cast_to(types, types, stream_types, False, __runtime__)
            ),
            __ctx__,
        )

    def SummarizeFileChunks(
        self,
        language: str,
        file_path: str,
        file_content: str,
        chunks: typing.List["types.ChunkInput"],
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlStream[stream_types.FileSummaryResult, types.FileSummaryResult]:
        __ctx__, __result__ = self.__options.merge_options(baml_options).create_async_stream(
            function_name="SummarizeFileChunks",
            args={
                "language": language,
                "file_path": file_path,
                "file_content": file_content,
                "chunks": chunks,
            },
        )
        return baml_py.BamlStream[stream_types.FileSummaryResult, types.FileSummaryResult](
            __result__,
            lambda x: typing.cast(
                stream_types.FileSummaryResult,
                x.cast_to(types, types, stream_types, True, __runtime__),
            ),
            lambda x: typing.cast(
                types.FileSummaryResult, x.cast_to(types, types, stream_types, False, __runtime__)
            ),
            __ctx__,
        )


class BamlHttpRequestClient:
    __options: DoNotUseDirectlyCallManager

    def __init__(self, options: DoNotUseDirectlyCallManager):
        self.__options = options

    async def SummarizeChunkAlone(
        self,
        language: str,
        symbol: str,
        line_start: int,
        line_end: int,
        content: str,
        parent_context: str,
        file_context: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        __result__ = await self.__options.merge_options(baml_options).create_http_request_async(
            function_name="SummarizeChunkAlone",
            args={
                "language": language,
                "symbol": symbol,
                "line_start": line_start,
                "line_end": line_end,
                "content": content,
                "parent_context": parent_context,
                "file_context": file_context,
            },
            mode="request",
        )
        return __result__

    async def SummarizeFileChunks(
        self,
        language: str,
        file_path: str,
        file_content: str,
        chunks: typing.List["types.ChunkInput"],
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        __result__ = await self.__options.merge_options(baml_options).create_http_request_async(
            function_name="SummarizeFileChunks",
            args={
                "language": language,
                "file_path": file_path,
                "file_content": file_content,
                "chunks": chunks,
            },
            mode="request",
        )
        return __result__


class BamlHttpStreamRequestClient:
    __options: DoNotUseDirectlyCallManager

    def __init__(self, options: DoNotUseDirectlyCallManager):
        self.__options = options

    async def SummarizeChunkAlone(
        self,
        language: str,
        symbol: str,
        line_start: int,
        line_end: int,
        content: str,
        parent_context: str,
        file_context: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        __result__ = await self.__options.merge_options(baml_options).create_http_request_async(
            function_name="SummarizeChunkAlone",
            args={
                "language": language,
                "symbol": symbol,
                "line_start": line_start,
                "line_end": line_end,
                "content": content,
                "parent_context": parent_context,
                "file_context": file_context,
            },
            mode="stream",
        )
        return __result__

    async def SummarizeFileChunks(
        self,
        language: str,
        file_path: str,
        file_content: str,
        chunks: typing.List["types.ChunkInput"],
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        __result__ = await self.__options.merge_options(baml_options).create_http_request_async(
            function_name="SummarizeFileChunks",
            args={
                "language": language,
                "file_path": file_path,
                "file_content": file_content,
                "chunks": chunks,
            },
            mode="stream",
        )
        return __result__


b = BamlAsyncClient(DoNotUseDirectlyCallManager({}))
