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 marketing_retriever import create_marketing_retriever

marketing_template = """You are support assitant of QWin online casino.
Your job is to answer questions about the marketing events and features in QWIN online casino, 
bonuses, user loyalty.
Use the following context to answer questions but don't mention it in your answer.
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 MarketingChainCreator:
    def __init__(self, 
                 llm: BaseChatModel):
        
        self.llm = llm
        
    def marketing_chain(self, 
                        user_id: str):
        retriever = create_marketing_retriever()
        
        docs_chain_prompt= ChatPromptTemplate.from_messages(
            [
                ('system', marketing_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