from typing import Dict, Any

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

from games_retriever import create_games_retriever

request_template = '''Given chat history, recognize the name of the game according to the list of games below: 

{game_list}

If a game is found in the game list, given the chat history, 
determine whether the user needs help or any details, result about a game already played 
or is simply asking about game rules, payments, and gameplay:

    1. If a user asks about a specific played game, its result or outcome reasons, determine the game id value: 
        - ONLY if the user has clearly indicated that this was the LAST game played, then define game id = "last".
        - otherwise try to determine game, numeric long parameter. User can find it on QWIN web site in "History" section
        - if there is no any game id indication - <empty string>
    
    2. if the user is asking about something about the game rules, game payouts, gameplay, 
    formulate a number of unique keywords for use when searching game information documents in the Qdrant document database. 
    Translate all keywords into the ENGLISH language, DO NOT use other languages.
    DO NOT include game name into the key words.
    Do NOT answer the question.

If the user asks to find games from all by a certain feature (for example: a certain symbol, slot theme, name, a certain payout, etc.):

    1. Try to determine search keywords for use when searching game information documents in the Qdrant document database. 
    
Generate and return only json string with format:
{{"game_codes":[<all recognized game codes>], 
  "game_names":[<all recognized game names>], 
  "key_words":[<list of key words translated into ENGLISH language>],
  "game_played":<is the user asking about a specific game, its result or outcome reasons? Values: true | false>,
  "game_id":<previously determined game id>,
  "search":<search keywords if user try to search games>}}

Human last message: {input}

Chat history:
'''

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

Human last message: {input}

Chat history:
"""

class GameChainCreator:
    def __init__(self, 
                 llm: BaseChatModel, 
                 initial_params: Dict[str, Any]):
        
        self.llm = llm
        self.initial_params = initial_params
        
        self.request_template = request_template.replace('{game_list}', initial_params['game_list'])
        
    def games_chain(self, 
                    user_id: str):
        
        question_prompt = ChatPromptTemplate.from_messages(
            [
                ('system', self.request_template),
                MessagesPlaceholder('chat_history'),
                ('assistant', 'JSON string:')
            ]
        )

        docs_chain_prompt= ChatPromptTemplate.from_messages(
            [
                ('system', games_template),
                MessagesPlaceholder('chat_history'),
                ('assistant', '')
            ]
        )

        retriever = create_games_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)
        games_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)

        return games_chain