from langchain_core.messages import BaseMessage
from langchain_core.prompts import PromptTemplate

from langflow.custom import CustomComponent
from langflow.field_typing import LanguageModel, Text


class ShouldRunNextComponent(CustomComponent):
    display_name = "Should Run Next"
    description = "Determines if a vertex is runnable."
    name = "ShouldRunNext"

    def build(self, llm: LanguageModel, question: str, context: str, retries: int = 3) -> Text:
        template = (
            "Given the following question and the context below, answer with a yes or no.\n\n"
            "{error_message}\n\n"
            "Question: {question}\n\n"  # noqa: RUF100, RUF027
            "Context: {context}\n\n"  # noqa: RUF100, RUF027
            "Answer:"
        )

        prompt = PromptTemplate.from_template(template)
        chain = prompt | llm
        error_message = ""
        for _i in range(retries):
            result = chain.invoke(
                {"question": question, "context": context, "error_message": error_message},
                config={"callbacks": self.get_langchain_callbacks()},
            )
            if isinstance(result, BaseMessage):
                content = result.content
            elif isinstance(result, str):
                content = result
            if isinstance(content, str) and content.lower().strip() in {"yes", "no"}:
                break
        condition = str(content).lower().strip() == "yes"
        self.status = f"Should Run Next: {condition}"
        if condition is False:
            self.stop()
        return context