"""Module for LLM answer generation chain."""

from typing import Any, Optional

from langchain_core.runnables import Runnable, RunnableConfig
from langchain_core.runnables.utils import Input
import inject

from langchain_ocr_lib.chains.chain import Chain
from langchain_ocr_lib.di_binding_keys.binding_keys import LangfuseManagerKey

RunnableInput = Input  # TODO: adjust properly
RunnableOutput = str


class OcrChain(Chain[RunnableInput, RunnableOutput]):
    """Base class for LLM answer generation chain."""

    _langfuse_manager = inject.attr(LangfuseManagerKey)

    def __init__(self):
        """Initialize the AnswerGenerationChain.

        Parameters
        ----------
        langfuse_manager : LangfuseManager
            Manager instance for handling Langfuse operations and monitoring
        """

    async def ainvoke(
        self, chain_input: RunnableInput, config: Optional[RunnableConfig] = None, **kwargs: Any
    ) -> RunnableOutput:
        """
        Asynchronously invokes the chain with given input.

        Parameters
        ----------
        chain_input : RunnableInput
            The input to be processed by the chain.
        chain_config : Optional[RunnableConfig]
            Configuration for the chain execution (default None).
        **kwargs : Any
            Additional keyword arguments passed to the chain.

        Returns
        -------
        RunnableOutput
            The output generated by the chain.

        Raises
        ------
        ChainError
            If an error occurs during chain execution.
        """
        return await self._create_chain().ainvoke(chain_input, config=config)

    def invoke(
        self, chain_input: RunnableInput, config: Optional[RunnableConfig] = None, **kwargs: Any
    ) -> RunnableOutput:
        """
        Invoke the chain with given input.

        Parameters
        ----------
        chain_input : RunnableInput
            The input to be processed by the chain.
        chain_config : Optional[RunnableConfig]
            Configuration for the chain execution (default None).
        **kwargs : Any
            Additional keyword arguments passed to the chain.

        Returns
        -------
        RunnableOutput
            The output generated by the chain.

        Raises
        ------
        ChainError
            If an error occurs during chain execution.
        """
        return self._create_chain().invoke(chain_input, config=config)

    def _create_chain(self) -> Runnable:
        return self._langfuse_manager.get_base_prompt(self.__class__.__name__) | self._langfuse_manager.get_base_llm(
            self.__class__.__name__
        )
