from promplate.chain import *#

class BaseCallback(Protocol):
    def pre_process(self, context: "ChainContext") -> Context | Awaitable[Context | None] | None: ...

    def mid_process(self, context: "ChainContext") -> Context | Awaitable[Context | None] | None: ...

    def end_process(self, context: "ChainContext") -> Context | Awaitable[Context | None] | None: ...

    def on_enter(self, node: "Interruptible", context: Context | None, config: Context) -> tuple[Context | None, Context]:
        return context, config

    def on_leave(self, node: "Interruptible", context: "ChainContext", config: Context) -> tuple["ChainContext", Context]:
        return context, config

class Callback(BaseCallback):
    def __init__(
        self,
        *,
        pre_process: "Process | AsyncProcess | None" = None,
        mid_process: "Process | AsyncProcess | None" = None,
        end_process: "Process | AsyncProcess | None" = None,
        on_enter: Callable[["Interruptible", Context | None, Context], tuple[Context | None, Context]] | None = None,
        on_leave: Callable[["Interruptible", "ChainContext", Context], tuple["ChainContext", Context]] | None = None,
    ):
        self._pre_process = pre_process
        self._mid_process = mid_process
        self._end_process = end_process
        self._on_enter = on_enter
        self._on_leave = on_leave

    def pre_process(self, context):
        if self._pre_process is not None:
            return self._pre_process(context)

    def mid_process(self, context):
        if self._mid_process is not None:
            return self._mid_process(context)

    def end_process(self, context):
        if self._end_process is not None:
            return self._end_process(context)

    def on_enter(self, node, context, config):
        if self._on_enter is not None:
            return self._on_enter(node, context, config)
        return context, config

    def on_leave(self, node, context, config):
        if self._on_leave is not None:
            return self._on_leave(node, context, config)
        return context, config

class Configurable:
    def __init__(self, **config):
        for key, val in config.items():
            setattr(self, key, val)

    @property
    def _config(self):
        return {k: v for k, v in self.__dict__.items() if not k.startswith("_")}

class Complete(Protocol):
    def __call__(self, prompt, /, **config) -> str: ...

class Generate(Protocol):
    def __call__(self, prompt, /, **config) -> Iterable[str]: ...

class AsyncComplete(Protocol):
    def __call__(self, prompt, /, **config) -> Awaitable[str]: ...

class AsyncGenerate(Protocol):
    def __call__(self, prompt, /, **config) -> AsyncIterable[str]: ...

class LLM(Protocol):
    @partial(cast, Complete | AsyncComplete)
    def complete(self, prompt, /, **config) -> str | Awaitable[str]: ...

    @partial(cast, Generate | AsyncGenerate)
    def generate(self, prompt, /, **config) -> Iterable[str] | AsyncIterable[str]: ...

class ChainContext(SafeChainMapContext):
    @overload
    def __new__(cls): ...

    @overload
    def __new__(cls, least: C, *maps: Mapping) -> C: ...

    @overload
    def __new__(cls, least: MutableMapping | None = None, *maps: Mapping): ...

    def __init__(self, least: MutableMapping | None = None, *maps: Mapping):
        super().__init__({} if least is None else least, *maps)  # type: ignore

    def __new__(cls, *args, **kwargs):  # type: ignore
        try:
            least = args[0]
        except IndexError:
            least = kwargs.get("least")
        if isinstance(least, cls) and least.__class__ is not cls:
            return least.__class__(*args, **kwargs)

        return super().__new__(cls, *args, **kwargs)

    @classmethod
    def ensure(cls, context):
        return context if isinstance(context, cls) else cls(context)

    @property
    def result(self):
        return self.__getitem__("__result__")

    @result.setter
    def result(self, result):
        self.__setitem__("__result__", result)

    @result.deleter
    def result(self):
        self.__delitem__("__result__")

    def __str__(self):
        return str({**self})

class AbstractNode(Protocol):
    def invoke(
        self,
        context: Context | None = None,
        /,
        complete: Complete | None = None,
        **config,
    ) -> ChainContext: ...

    async def ainvoke(
        self,
        context: Context | None = None,
        /,
        complete: Complete | AsyncComplete | None = None,
        **config,
    ) -> ChainContext: ...

    def stream(
        self,
        context: Context | None = None,
        /,
        generate: Generate | None = None,
        **config,
    ) -> Iterable[ChainContext]: ...

    def astream(
        self,
        context: Context | None = None,
        /,
        generate: Generate | AsyncGenerate | None = None,
        **config,
    ) -> AsyncIterable[ChainContext]: ...

    @classmethod
    def _get_chain_type(cls):
        return Chain

    def __add__(self, chain: "AbstractNode"):
        if isinstance(chain, Chain):
            return self._get_chain_type()(self, *chain)
        return self._get_chain_type()(self, chain)

class Interruptible(AbstractNode, Protocol):
    def _invoke(
        self,
        context: ChainContext,
        /,
        complete: Complete | None,
        callbacks: list[BaseCallback],
        **config,
    ): ...

    async def _ainvoke(
        self,
        context: ChainContext,
        /,
        complete: Complete | AsyncComplete | None,
        callbacks: list[BaseCallback],
        **config,
    ): ...

    def _stream(
        self,
        context: ChainContext,
        /,
        generate: Generate | None,
        callbacks: list[BaseCallback],
        **config,
    ) -> Iterable: ...

    def _astream(
        self,
        context: ChainContext,
        /,
        generate: Generate | AsyncGenerate | None,
        callbacks: list[BaseCallback],
        **config,
    ) -> AsyncIterable: ...

    callbacks: list[BaseCallback | type[BaseCallback]]

    def enter(self, context: Context | None, config: Context):
        callbacks: list[BaseCallback] = ensure_callbacks(self.callbacks)
        for callback in callbacks:
            context, config = callback.on_enter(self, context, config)
        return context, config, callbacks

    def leave(self, context: ChainContext, config: Context, callbacks: list[BaseCallback]):
        for callback in reversed(callbacks):
            context, config = callback.on_leave(self, context, config)
        return context, config

    def add_pre_processes(self, *processes: Process | AsyncProcess):
        self.callbacks.extend(Callback(pre_process=i) for i in processes)
        return self

    def add_mid_processes(self, *processes: Process | AsyncProcess):
        self.callbacks.extend(Callback(mid_process=i) for i in processes)
        return self

    def add_end_processes(self, *processes: Process | AsyncProcess):
        self.callbacks.extend(Callback(end_process=i) for i in processes)
        return self

    def add_callbacks(self, *callbacks: BaseCallback | type[BaseCallback]):
        self.callbacks.extend(callbacks)
        return self

    def pre_process(self, process: Process | AsyncProcess):
        self.add_pre_processes(process)
        return process

    def mid_process(self, process: Process | AsyncProcess):
        self.add_mid_processes(process)
        return process

    def end_process(self, process: Process | AsyncProcess):
        self.add_end_processes(process)
        return process

    def callback(self, callback: BaseCallback | type[BaseCallback]):
        self.add_callbacks(callback)
        return callback

    @staticmethod
    def _apply_pre_processes(context: ChainContext, callbacks: list[BaseCallback]):
        for callback in callbacks:
            context |= cast(Context, callback.pre_process(context) or {})

    @staticmethod
    def _apply_mid_processes(context: ChainContext, callbacks: list[BaseCallback]):
        for callback in callbacks:
            context |= cast(Context, callback.mid_process(context) or {})

    @staticmethod
    def _apply_end_processes(context: ChainContext, callbacks: list[BaseCallback]):
        for callback in reversed(callbacks):
            context |= cast(Context, callback.end_process(context) or {})

    @staticmethod
    async def _apply_async_pre_processes(context: ChainContext, callbacks: list[BaseCallback]):
        for callback in callbacks:
            context |= cast(Context, await resolve(callback.pre_process(context)) or {})

    @staticmethod
    async def _apply_async_mid_processes(context: ChainContext, callbacks: list[BaseCallback]):
        for callback in callbacks:
            context |= cast(Context, await resolve(callback.mid_process(context)) or {})

    @staticmethod
    async def _apply_async_end_processes(context: ChainContext, callbacks: list[BaseCallback]):
        for callback in reversed(callbacks):
            context |= cast(Context, await resolve(callback.end_process(context)) or {})

    def invoke(self, context=None, /, complete=None, **config) -> ChainContext:
        context, config, callbacks = self.enter(context, config)
        context = ChainContext.ensure(context)

        try:
            self._invoke(ChainContext(context, self.context), complete, callbacks, **config)
        except Jump as jump:
            context, config = self.leave(context, config, callbacks)
            if jump.out_of is not None and jump.out_of is not self:
                raise jump from None
            if jump.into is not None:
                jump.into.invoke(context, complete, **config)
        else:
            context, config = self.leave(context, config, callbacks)

        return context

    async def ainvoke(self, context=None, /, complete=None, **config) -> ChainContext:
        context, config, callbacks = self.enter(context, config)
        context = ChainContext.ensure(context)

        try:
            await self._ainvoke(ChainContext(context, self.context), complete, callbacks, **config)
        except Jump as jump:
            context, config = self.leave(context, config, callbacks)
            if jump.out_of is not None and jump.out_of is not self:
                raise jump from None
            if jump.into is not None:
                await jump.into.ainvoke(context, complete, **config)
        else:
            context, config = self.leave(context, config, callbacks)

        return context

    def stream(self, context=None, /, generate=None, **config) -> Iterable[ChainContext]:
        context, config, callbacks = self.enter(context, config)
        context = ChainContext.ensure(context)

        try:
            for _ in self._stream(ChainContext(context, self.context), generate, callbacks, **config):
                yield context
        except Jump as jump:
            context, config = self.leave(context, config, callbacks)
            if jump.out_of is not None and jump.out_of is not self:
                raise jump from None
            if jump.into is not None:
                yield from jump.into.stream(context, generate, **config)
        else:
            context, config = self.leave(context, config, callbacks)

    async def astream(self, context=None, /, generate=None, **config) -> AsyncIterable[ChainContext]:
        context, config, callbacks = self.enter(context, config)
        context = ChainContext.ensure(context)

        try:
            async for _ in self._astream(ChainContext(context, self.context), generate, callbacks, **config):
                yield context
        except Jump as jump:
            context, config = self.leave(context, config, callbacks)
            if jump.out_of is not None and jump.out_of is not self:
                raise jump from None
            if jump.into is not None:
                async for i in jump.into.astream(context, generate, **config):
                    yield i
        else:
            context, config = self.leave(context, config, callbacks)

    _context: Context | None

    @property
    def context(self):
        if self._context is None:
            self._context = {}
        return self._context

    @context.setter
    def context(self, context: Context | None):
        self._context = context

    @context.deleter
    def context(self):
        self._context = None