import json
import os
import pathlib
from typing import Dict, List, Tuple

import weaviate
from langchain import OpenAI, PromptTemplate
from langchain.chains import LLMChain
from langchain.chains.base import Chain
from langchain.chains.combine_documents.base import BaseCombineDocumentsChain
from langchain.chains.conversation.memory import ConversationBufferMemory

from langchain.chains.question_answering import load_qa_chain #BaseCombineDocumentsChain
from langchain.embeddings import OpenAIEmbeddings
from langchain.prompts import FewShotPromptTemplate, PromptTemplate
from langchain.prompts.example_selector import \
    SemanticSimilarityExampleSelector
from langchain.vectorstores import FAISS, Weaviate
from pydantic import BaseModel

from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import FAISS, Weaviate
os.environ["WEAVIATE_URL"] = "https://tro.weaviate.network/"
os.environ["OPENAI_API_KEY"] = "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"
KEY ="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"

class CustomChain(Chain, BaseModel):

    vstore: Weaviate
    chain: BaseCombineDocumentsChain
    key_word_extractor: Chain

    @property
    def input_keys(self) -> List[str]:
        return ["game_description"]

    @property
    def output_keys(self) -> List[str]:
        return ["game_environment"]

    def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:

        game_description = inputs["game_description"]
        chat_history_str = _get_chat_history(inputs["chat_history"])

        if chat_history_str: # Takes in user input
            print('running key_word_extractor')
            new_game_description= self.key_word_extractor.run(
                game_description=game_description, chat_history=chat_history_str
            )
            print('new_game_description',new_game_description)
        else:
            new_game_description = game_description
            print('_new_game_description', new_game_description) # sendig objects , not showing up on screen ? 

        docs = self.vstore.similarity_search(new_game_description, k=4)
        
        new_inputs = inputs.copy()
        new_inputs["game_description"] = new_game_description
        new_inputs["chat_history"] = chat_history_str

        game_environment, _ = self.chain.combine_docs(docs, **new_inputs) # StuffDocumentsChain
        print('combined', game_environment) # doc combo of game descripton and chat history,

        return {"game_environment": game_environment}


# 1. Given the chat history and new user input, 
# Generate game environment, objects, npcs, enemies, and quest narrative.

def get_new_chain1(vectorstore) -> Chain:
    # Calls Vector DB 
    WEAVIATE_URL = os.environ["WEAVIATE_URL"]
    client = weaviate.Client(
        url=WEAVIATE_URL,
        additional_headers={"X-OpenAI-Api-Key": os.environ["OPENAI_API_KEY"]},
    )
    _eg_template = """## AI:

    Chat History:
    {chat_history}
    Follow Up Input: {game_description} 
    standalone detail: {game_description} {game_environment}"""
    _eg_prompt = PromptTemplate(
        template=_eg_template,
        input_variables=["chat_history","game_description", "game_environment"], # question aka game detail ################ showing up in Terminal? ##############
    )

    # _prefix = """Given the following conversation and a follow up game detail, summarize the follow up game detail 
    #in a way that is coherent with the conversation.\
    #  You should assume that the detail is related to Game Design. Never change the game Title unless the human asks you to."""
    _prefix = """ If the input and the follow up input are closely related, and are both related to the game suggested by 
    the previous conversation and if you are completely confident if it the same game from the original converation, then proceed. Otherwise, clarify the game title and proceed.
    To proceed, retrieve the game Title and share again in quotes to the human, and summarize the follow up input in a way that is coherent with the conversation. 
    If the input detail is related to Game Design, proceed, otherwise clarify.
        """

    _suffix = """## AI:

    Chat History:
    {chat_history}
    Follow Up Input: {game_description}
    standalone detail:"""
    eg_store = Weaviate(
        client,
        "Rephrase",
        "content",
        attributes=["game_description", "game_environment", "chat_history"], 
    )

    example_selector = SemanticSimilarityExampleSelector(vectorstore=eg_store, k=4) # looks for similar examples in vector db
    prompt = FewShotPromptTemplate(
        prefix=_prefix, # prompt
        suffix=_suffix,
        example_selector=example_selector,
        example_prompt=_eg_prompt,
        input_variables=["game_description", "chat_history"],
    )
    llm = OpenAI(temperature=0.7, model_name="text-davinci-003", max_tokens=1000)
    key_word_extractor = LLMChain(llm=llm, prompt=prompt) # this could be game word extractor

    # Gives an example and relevant documents # from Vector DB  ##################################################### showing up in UI ##############
    EXAMPLE_PROMPT = PromptTemplate(
        template=">AI:\nContent:\n---------\n{page_content}\n----------\n", # not working right now ? 
        input_variables=["page_content"],
    )

    template = """
    You are an AI assistant for generating a single game, the central narrative and its game quests inside the same game, 
    conversing in a succinct, coherent and interactive manner. Assume the entire conversation is about a single game only.

    You are given the game description. You are also given the chat history, which is a conversation between a human and 
    another AI assistant about the same game. Do not answer any questions from the other AI assistant. 
    If you don't know the answer, just say "I'm not sure, how this is related?" and ask the human to clarify.

    If asked about a quest, assume it is within the same game, 
    and provide a quest motivation, participants, rewards, character dialogue relevant to the game.

    From the input, introduce a short back story of the game's plot points or character motivations and character dialogue.
    Plot points may include an inciting incidient, a rising action composed of a series of conflicts, complications, dilemmas, obstacles.

    From the input, generate game objects (collectibles, NPCs, enemies) related to the plot points, limit to 1-2 objects.
    Ask a short question to clarify details if you are unsure but you are an expert so offer a suggestion and only ask once.
    
    If the human asks for or adds game objects, provide only 1 detail such as physical description and possible player interactions with the object.
    If the human asks for or adds an enemy, describe its behaviors and personality, as related to the game plot points, limit to a few words.
    If the human asks for or adds a collectible, describe its physical appearance and possible player interactions with the object.
    If the human asks for or adds a noun or NPC, describe its behaviors and personality, as related to the game plot points, limit to a few words.
    If the huamn asks for character dialogue, provide a short dialogue snippet, with 1-2 sentences, as related to the game plot points.
    Limit your reponses to 3-6 sentences.
    
    Game Description: {game_description}
    {context}
    Answer in Markdown:"""

    PROMPT = PromptTemplate(template=template, input_variables=["game_description", "context"])
    doc_chain = load_qa_chain(
        OpenAI(temperature=0.7, model_name="text-davinci-003", max_tokens=1000),
        chain_type="stuff", # all related data is stuffed into prompt as context
        prompt=PROMPT,
        document_prompt=EXAMPLE_PROMPT,
    )
    return CustomChain(
        chain=doc_chain, vstore=vectorstore, key_word_extractor=key_word_extractor
    )
    # doc + input -> game environment

def _get_chat_history(chat_history: List[Tuple[str, str]]):
    buffer = ""
    for human_s, ai_s in chat_history:
        human = f"Human: " + human_s
        ai = f"Assistant: " + ai_s
        buffer += "\n" + "\n".join([human, ai])
    return buffer