from langchain.prompts import ChatPromptTemplate
from langchain_core.language_models.chat_models import BaseChatModel
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains import create_retrieval_chain

from affiliate_retriever import create_affiliate_retriever

affiliate_template = """You are support assitant of QWin online casino.
Your job is to answer questions about the affiliate program for QWIN online casino, 
promo links, affiliate users, affiliate payments and benefits.
Use the following context to answer questions.
Be as detailed as possible, but don't make up any information
that's not from the context. If you don't know an answer,
say you don't know.
{context}
"""

class AffiliateChainCreator:
    def __init__(self, 
                 llm: BaseChatModel):
        
        self.llm = llm
        
    def affiliate_chain(self, 
                        user_id: str):
        retriever = create_affiliate_retriever()
        
        docs_chain_prompt= ChatPromptTemplate.from_messages(
            [
                ('system', affiliate_template),
                ('assistant', '')
            ]
        )
        
        question_answer_chain = create_stuff_documents_chain(self.llm, docs_chain_prompt)
        affiliate_chain = create_retrieval_chain(retriever, question_answer_chain)
        
        return affiliate_chain