###############################################################################
#
#  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.
#
# ruff: noqa: E501,F401
# flake8: noqa: E501,F401
# pylint: disable=unused-import,line-too-long
# fmt: off
from typing import Any, Dict, List, Optional, TypeVar, Union, TypedDict, Type
from typing_extensions import NotRequired
import pprint

import baml_py
from pydantic import BaseModel, ValidationError, create_model

from . import partial_types, types
from .type_builder import TypeBuilder
from .globals import DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX, DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME


OutputType = TypeVar('OutputType')

def coerce(cls: Type[BaseModel], parsed: Any) -> Any:
  try:
    return cls.model_validate({"inner": parsed}).inner # type: ignore
  except ValidationError as e:
    raise TypeError(
      "Internal BAML error while casting output to {}\n{}".format(
        cls.__name__,
        pprint.pformat(parsed)
      )
    ) from e

# Define the TypedDict with optional parameters having default values
class BamlCallOptions(TypedDict, total=False):
    tb: NotRequired[TypeBuilder]
    client_registry: NotRequired[baml_py.baml_py.ClientRegistry]

class BamlAsyncClient:
    __runtime: baml_py.BamlRuntime
    __ctx_manager: baml_py.BamlCtxManager
    __stream_client: "BamlStreamClient"

    def __init__(self, runtime: baml_py.BamlRuntime, ctx_manager: baml_py.BamlCtxManager):
      self.__runtime = runtime
      self.__ctx_manager = ctx_manager
      self.__stream_client = BamlStreamClient(self.__runtime, self.__ctx_manager)

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


    
    async def ChangePlan(
        self,
        screenshot: baml_py.Image,changes: str,old_plan: List[types.PlanningStep],
        baml_options: BamlCallOptions = {},
    ) -> List[types.PlanningStep]:
      __tb__ = baml_options.get("tb", None)
      if __tb__ is not None:
        tb = __tb__._tb
      else:
        tb = None
      __cr__ = baml_options.get("client_registry", None)

      raw = await self.__runtime.call_function(
        "ChangePlan",
        {
          "screenshot": screenshot,"changes": changes,"old_plan": old_plan,
        },
        self.__ctx_manager.get(),
        tb,
        __cr__,
      )
      mdl = create_model("ChangePlanReturnType", inner=(List[types.PlanningStep], ...))
      return coerce(mdl, raw.parsed())
    
    async def CreatePlan(
        self,
        screenshot: baml_py.Image,task_instructions: str,old_plan: List[types.PlanningStep],
        baml_options: BamlCallOptions = {},
    ) -> List[types.PlanningStep]:
      __tb__ = baml_options.get("tb", None)
      if __tb__ is not None:
        tb = __tb__._tb
      else:
        tb = None
      __cr__ = baml_options.get("client_registry", None)

      raw = await self.__runtime.call_function(
        "CreatePlan",
        {
          "screenshot": screenshot,"task_instructions": task_instructions,"old_plan": old_plan,
        },
        self.__ctx_manager.get(),
        tb,
        __cr__,
      )
      mdl = create_model("CreatePlanReturnType", inner=(List[types.PlanningStep], ...))
      return coerce(mdl, raw.parsed())
    
    async def DetermineSelectorStyle(
        self,
        instructions: str,target: types.Element,next_sibling: Optional[types.Element],previous_sibling: Optional[types.Element],ancestors: str,
        baml_options: BamlCallOptions = {},
    ) -> List[types.SelectorStyle]:
      __tb__ = baml_options.get("tb", None)
      if __tb__ is not None:
        tb = __tb__._tb
      else:
        tb = None
      __cr__ = baml_options.get("client_registry", None)

      raw = await self.__runtime.call_function(
        "DetermineSelectorStyle",
        {
          "instructions": instructions,"target": target,"next_sibling": next_sibling,"previous_sibling": previous_sibling,"ancestors": ancestors,
        },
        self.__ctx_manager.get(),
        tb,
        __cr__,
      )
      mdl = create_model("DetermineSelectorStyleReturnType", inner=(List[types.SelectorStyle], ...))
      return coerce(mdl, raw.parsed())
    
    async def ElementsPerStep(
        self,
        plan: str,current_step: types.PlanningStep,
        baml_options: BamlCallOptions = {},
    ) -> List[str]:
      __tb__ = baml_options.get("tb", None)
      if __tb__ is not None:
        tb = __tb__._tb
      else:
        tb = None
      __cr__ = baml_options.get("client_registry", None)

      raw = await self.__runtime.call_function(
        "ElementsPerStep",
        {
          "plan": plan,"current_step": current_step,
        },
        self.__ctx_manager.get(),
        tb,
        __cr__,
      )
      mdl = create_model("ElementsPerStepReturnType", inner=(List[str], ...))
      return coerce(mdl, raw.parsed())
    
    async def GeneratePlaywrightCode(
        self,
        procedure_name: str,procedure_steps: List[types.Step],
        baml_options: BamlCallOptions = {},
    ) -> str:
      __tb__ = baml_options.get("tb", None)
      if __tb__ is not None:
        tb = __tb__._tb
      else:
        tb = None
      __cr__ = baml_options.get("client_registry", None)

      raw = await self.__runtime.call_function(
        "GeneratePlaywrightCode",
        {
          "procedure_name": procedure_name,"procedure_steps": procedure_steps,
        },
        self.__ctx_manager.get(),
        tb,
        __cr__,
      )
      mdl = create_model("GeneratePlaywrightCodeReturnType", inner=(str, ...))
      return coerce(mdl, raw.parsed())
    
    async def GenerateXPathSelector(
        self,
        instructions: str,target: types.Element,ancestors: str,next_sibling: Optional[types.Element],previous_sibling: Optional[types.Element],selector_styles: List[types.SelectorStyle],
        baml_options: BamlCallOptions = {},
    ) -> List[str]:
      __tb__ = baml_options.get("tb", None)
      if __tb__ is not None:
        tb = __tb__._tb
      else:
        tb = None
      __cr__ = baml_options.get("client_registry", None)

      raw = await self.__runtime.call_function(
        "GenerateXPathSelector",
        {
          "instructions": instructions,"target": target,"ancestors": ancestors,"next_sibling": next_sibling,"previous_sibling": previous_sibling,"selector_styles": selector_styles,
        },
        self.__ctx_manager.get(),
        tb,
        __cr__,
      )
      mdl = create_model("GenerateXPathSelectorReturnType", inner=(List[str], ...))
      return coerce(mdl, raw.parsed())
    


class BamlStreamClient:
    __runtime: baml_py.BamlRuntime
    __ctx_manager: baml_py.BamlCtxManager

    def __init__(self, runtime: baml_py.BamlRuntime, ctx_manager: baml_py.BamlCtxManager):
      self.__runtime = runtime
      self.__ctx_manager = ctx_manager

    
    def ChangePlan(
        self,
        screenshot: baml_py.Image,changes: str,old_plan: List[types.PlanningStep],
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlStream[List[partial_types.PlanningStep], List[types.PlanningStep]]:
      __tb__ = baml_options.get("tb", None)
      if __tb__ is not None:
        tb = __tb__._tb
      else:
        tb = None
      __cr__ = baml_options.get("client_registry", None)

      raw = self.__runtime.stream_function(
        "ChangePlan",
        {
          "screenshot": screenshot,
          "changes": changes,
          "old_plan": old_plan,
        },
        None,
        self.__ctx_manager.get(),
        tb,
        __cr__,
      )

      mdl = create_model("ChangePlanReturnType", inner=(List[types.PlanningStep], ...))
      partial_mdl = create_model("ChangePlanPartialReturnType", inner=(List[partial_types.PlanningStep], ...))

      return baml_py.BamlStream[List[partial_types.PlanningStep], List[types.PlanningStep]](
        raw,
        lambda x: coerce(partial_mdl, x),
        lambda x: coerce(mdl, x),
        self.__ctx_manager.get(),
      )
    
    def CreatePlan(
        self,
        screenshot: baml_py.Image,task_instructions: str,old_plan: List[types.PlanningStep],
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlStream[List[partial_types.PlanningStep], List[types.PlanningStep]]:
      __tb__ = baml_options.get("tb", None)
      if __tb__ is not None:
        tb = __tb__._tb
      else:
        tb = None
      __cr__ = baml_options.get("client_registry", None)

      raw = self.__runtime.stream_function(
        "CreatePlan",
        {
          "screenshot": screenshot,
          "task_instructions": task_instructions,
          "old_plan": old_plan,
        },
        None,
        self.__ctx_manager.get(),
        tb,
        __cr__,
      )

      mdl = create_model("CreatePlanReturnType", inner=(List[types.PlanningStep], ...))
      partial_mdl = create_model("CreatePlanPartialReturnType", inner=(List[partial_types.PlanningStep], ...))

      return baml_py.BamlStream[List[partial_types.PlanningStep], List[types.PlanningStep]](
        raw,
        lambda x: coerce(partial_mdl, x),
        lambda x: coerce(mdl, x),
        self.__ctx_manager.get(),
      )
    
    def DetermineSelectorStyle(
        self,
        instructions: str,target: types.Element,next_sibling: Optional[types.Element],previous_sibling: Optional[types.Element],ancestors: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlStream[List[Optional[types.SelectorStyle]], List[types.SelectorStyle]]:
      __tb__ = baml_options.get("tb", None)
      if __tb__ is not None:
        tb = __tb__._tb
      else:
        tb = None
      __cr__ = baml_options.get("client_registry", None)

      raw = self.__runtime.stream_function(
        "DetermineSelectorStyle",
        {
          "instructions": instructions,
          "target": target,
          "next_sibling": next_sibling,
          "previous_sibling": previous_sibling,
          "ancestors": ancestors,
        },
        None,
        self.__ctx_manager.get(),
        tb,
        __cr__,
      )

      mdl = create_model("DetermineSelectorStyleReturnType", inner=(List[types.SelectorStyle], ...))
      partial_mdl = create_model("DetermineSelectorStylePartialReturnType", inner=(List[Optional[types.SelectorStyle]], ...))

      return baml_py.BamlStream[List[Optional[types.SelectorStyle]], List[types.SelectorStyle]](
        raw,
        lambda x: coerce(partial_mdl, x),
        lambda x: coerce(mdl, x),
        self.__ctx_manager.get(),
      )
    
    def ElementsPerStep(
        self,
        plan: str,current_step: types.PlanningStep,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlStream[List[Optional[str]], List[str]]:
      __tb__ = baml_options.get("tb", None)
      if __tb__ is not None:
        tb = __tb__._tb
      else:
        tb = None
      __cr__ = baml_options.get("client_registry", None)

      raw = self.__runtime.stream_function(
        "ElementsPerStep",
        {
          "plan": plan,
          "current_step": current_step,
        },
        None,
        self.__ctx_manager.get(),
        tb,
        __cr__,
      )

      mdl = create_model("ElementsPerStepReturnType", inner=(List[str], ...))
      partial_mdl = create_model("ElementsPerStepPartialReturnType", inner=(List[Optional[str]], ...))

      return baml_py.BamlStream[List[Optional[str]], List[str]](
        raw,
        lambda x: coerce(partial_mdl, x),
        lambda x: coerce(mdl, x),
        self.__ctx_manager.get(),
      )
    
    def GeneratePlaywrightCode(
        self,
        procedure_name: str,procedure_steps: List[types.Step],
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlStream[Optional[str], str]:
      __tb__ = baml_options.get("tb", None)
      if __tb__ is not None:
        tb = __tb__._tb
      else:
        tb = None
      __cr__ = baml_options.get("client_registry", None)

      raw = self.__runtime.stream_function(
        "GeneratePlaywrightCode",
        {
          "procedure_name": procedure_name,
          "procedure_steps": procedure_steps,
        },
        None,
        self.__ctx_manager.get(),
        tb,
        __cr__,
      )

      mdl = create_model("GeneratePlaywrightCodeReturnType", inner=(str, ...))
      partial_mdl = create_model("GeneratePlaywrightCodePartialReturnType", inner=(Optional[str], ...))

      return baml_py.BamlStream[Optional[str], str](
        raw,
        lambda x: coerce(partial_mdl, x),
        lambda x: coerce(mdl, x),
        self.__ctx_manager.get(),
      )
    
    def GenerateXPathSelector(
        self,
        instructions: str,target: types.Element,ancestors: str,next_sibling: Optional[types.Element],previous_sibling: Optional[types.Element],selector_styles: List[types.SelectorStyle],
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlStream[List[Optional[str]], List[str]]:
      __tb__ = baml_options.get("tb", None)
      if __tb__ is not None:
        tb = __tb__._tb
      else:
        tb = None
      __cr__ = baml_options.get("client_registry", None)

      raw = self.__runtime.stream_function(
        "GenerateXPathSelector",
        {
          "instructions": instructions,
          "target": target,
          "ancestors": ancestors,
          "next_sibling": next_sibling,
          "previous_sibling": previous_sibling,
          "selector_styles": selector_styles,
        },
        None,
        self.__ctx_manager.get(),
        tb,
        __cr__,
      )

      mdl = create_model("GenerateXPathSelectorReturnType", inner=(List[str], ...))
      partial_mdl = create_model("GenerateXPathSelectorPartialReturnType", inner=(List[Optional[str]], ...))

      return baml_py.BamlStream[List[Optional[str]], List[str]](
        raw,
        lambda x: coerce(partial_mdl, x),
        lambda x: coerce(mdl, x),
        self.__ctx_manager.get(),
      )
    

b = BamlAsyncClient(DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME, DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX)

__all__ = ["b"]