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 casino_retriever import create_casino_retriever

casino_template = """You are support assistant of QWIN online casino.
Your job is to answer questions about the operation of QWIN online casino, 
casino rules, casino restrictions, user registration, user account. 
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 CasinoChainCreator:
    def __init__(self, 
                 llm: BaseChatModel):
        
        self.llm = llm
        
    def casino_chain(self, 
                        user_id: str):
        retriever = create_casino_retriever()
        
        docs_chain_prompt= ChatPromptTemplate.from_messages(
            [
                ('system', casino_template),
                ('assistant', '')
            ]
        )
        
        question_answer_chain = create_stuff_documents_chain(self.llm, docs_chain_prompt)
        casino_chain = create_retrieval_chain(retriever, question_answer_chain)
        
        return casino_chain