from typing import Any, Optional
from langchain.chains import LLMChain
from langchain.base_language import BaseLanguageModel
from langchain.schema import LLMResult, PromptValue
from langchain.prompts import PromptTemplate
from langchain.memory.chat_memory import BaseMemory
from langchain.chat_models import ChatOpenAI

from promopts import CONTENT_RE_WRIGHT_PROMPT, FEEDBACK_PROMPT


class HumanFeedBackChain(LLMChain):
    """Chain to run queries against LLMs."""

    memory: Optional[BaseMemory] = None

    def __init__(self, verbose=True, llm: BaseLanguageModel = ChatOpenAI(temperature=0.7),  memory: Optional[BaseMemory] = None, prompt: PromptTemplate = FEEDBACK_PROMPT):
        super().__init__(llm=llm, prompt=prompt, memory=memory, verbose=verbose)

    def run(self, *args: Any, **kwargs: Any) -> str:
        """Run the chain as text in, text out or multiple variables, text out."""
        if len(self.output_keys) != 1:
            raise ValueError(
                f"`run` not supported when there is not exactly "
                f"one output key. Got {self.output_keys}."
            )

        if args and not kwargs:
            if len(args) != 1:
                raise ValueError(
                    "`run` supports only one positional argument.")
            return self("Answer:" + args[0])[self.output_keys[0]]

        if kwargs and not args:
            return self(kwargs)[self.output_keys[0]]

        raise ValueError(
            f"`run` supported with either positional arguments or keyword arguments"
            f" but not both. Got args: {args} and kwargs: {kwargs}."
        )


contextRewriteChain = LLMChain(llm=ChatOpenAI(temperature=0.7), prompt=CONTENT_RE_WRIGHT_PROMPT)