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

from finance_retriever import create_finance_retriever

request_template = '''Role: ou are support assistant of QWin online casino for financial queries.

Task: Based on the chat history, identify the category of the user's question: 
- "terms": General questions about terms and conditions, deposits, withdrawals, and restrictions.
- "deposit": User has issues with a specific deposit. Try to identify the Bitcoin transaction hash in the blockchain and the amount.
- "withdrawal": User has issues with a specific withdrawal. Try to identify the amount.
- "other": Everything else.

Generate and return only json string with format:
{{
    "category": <recognized category>,
    "data": {{
        "trx_hash": <hash of deposit transaction or empty string>,
        "amount": <amount of deposit or withdrawal in satoshi or empty string>
    }}
}}

Provide the category and details based on the chat history. Use the following keys for the JSON output:
- "category": One of "terms", "deposit", "withdrawal", "other".
- "trx_hash": Transaction hash if available, otherwise an empty string.
- "amount": Amount in satoshi if available, otherwise an empty string.

Human last message: {input}

Chat history:
'''

finance_template = """You are support assistant of QWin online casino.
Your job is to answer questions about the money operation of QWIN online casino, 
rules, restrictions, deposits, withdrawals. 
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:
{context}

Human last message: {input}

Chat history:
"""

class FinanceChainCreator:
    def __init__(self, 
                 llm: BaseChatModel):
        
        self.llm = llm
        
    def finance_chain(self, 
                      user_id: str):
        question_prompt = ChatPromptTemplate.from_messages(
            [
                ('system', request_template),
                MessagesPlaceholder('chat_history'),
                ('assistant', 'JSON string:')
            ]
        )
        
        docs_chain_prompt= ChatPromptTemplate.from_messages(
            [
                ('system', finance_template),
                MessagesPlaceholder('chat_history'),
                ('assistant', '')
            ]
        )
        
        retriever = create_finance_retriever(user_id=user_id)
        history_aware_retriever = create_history_aware_retriever(self.llm, retriever, question_prompt)
        
        question_answer_chain = create_stuff_documents_chain(self.llm, docs_chain_prompt)
        finance_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
        
        return finance_chain