import os

import utils

import traceback
from langchain.chains.qa_with_sources import load_qa_with_sources_chain
from langchain.chains import ConversationChain
from langchain.llms import OpenAI
import langchain
from langchain.cache import InMemoryCache
from langchain.llms import OpenAI
from langchain.chains.conversation.memory import ConversationSummaryBufferMemory,ConversationBufferMemory,ConversationBufferWindowMemory
from langchain.prompts import PromptTemplate
from embeddings import EmbeddingsManager
from flask import Flask, send_from_directory
import json
import time
import threading
import secrets
import string
import hashlib
from flask import request
from langchain.cache import InMemoryCache,SQLiteCache
import re
import requests
from waitress import serve
from translator import Translator
import sys
from query.discoursequery import DiscourseQuery
from query.embeddingsquery import EmbeddingsQuery
from Summary import Summary
import uuid
from langchain.llms import NLPCloud
from langchain.llms import AI21
from langchain.llms import Cohere
from SmartCache import SmartCache
CONFIG=None
QUERIERS=[]

args=sys.argv
confiFile=args[1] if len(args)>1 else "config.json"
print("Use config file", confiFile)
with open(confiFile, "r") as f:
    CONFIG=json.load(f)
    EmbeddingsManager.init(CONFIG)
    Summary.init(CONFIG)
    QUERIERS=[
        EmbeddingsQuery(CONFIG),
        DiscourseQuery(
            CONFIG,CONFIG["JME_HUB_URL"],
            searchFilter=CONFIG["JME_HUB_SEARCH_FILTER"],
            knowledgeCutoff=CONFIG["JME_HUB_KNOWLEDGE_CUTOFF"]
        )
    ]
    Translator.init(CONFIG)

def getAffineDocs(question,context,keywords,shortQuestion, wordSalad=None, unitFilter=None,
    maxFragmentsToReturn=3, maxFragmentsToSelect=12,merge=False):
    affineDocs=[]

   
    for q in QUERIERS:
        print("Get affine docs from",q,"using question",question,"with context",context,"and keywords",keywords)
        t=time.time()
        v=q.getAffineDocs(
            question, context, keywords,shortQuestion, wordSalad, unitFilter,
            maxFragmentsToReturn=maxFragmentsToReturn,
            maxFragmentsToSelect=maxFragmentsToSelect,
            merge=merge        
        )
        print("Completed in",time.time()-t,"seconds.")
        if v!=None:
            affineDocs.extend(v)
    return affineDocs
    
def rewriteError(error):
    if error.startswith("Rate limit reached ") :
        return "Rate limit."

def rewrite(question):
    # replace app, applet, game, application with simple application 
    question=re.sub(r"\b(app|applet|game|application)\b", "simple application", question, flags=re.IGNORECASE)

    return question



def createChain():
    



    # Backward compatibility
    model_name=CONFIG.get("OPENAI_MODEL","text-davinci-003")
    llm_name="openai"
    ######## 
    
    llmx=CONFIG.get("LLM_MODEL",None) # "openai:text-davinci-003" "cohere:xlarge"
    if llmx!=None: 
        if ":" in llmx:
            llm_name,model_name=llmx.split(":")
        else:
            llm_name,model_name=llmx.split(".")



    template = ""
    template_path="prompts/"+llm_name+"."+model_name+".txt"
    if not os.path.exists(template_path):
        template_path="prompts/openai.text-davinci-003.txt"
    
    with open(template_path, "r") as f:
        template=f.read()

    prompt = PromptTemplate(
        input_variables=[ "history", "question", "summaries"], 
        template=template
    )

    llm=None
    history_length=700
    if llm_name=="openai":
        max_tokens=512
        temperature=0.0
        if model_name=="text-davinci-003":
            max_tokens=512
        elif model_name=="code-davinci-002":
            max_tokens=1024
            #history_length=1024            
        llm=OpenAI(
            temperature=temperature,
            model_name=model_name,
            max_tokens=max_tokens,
        )
    elif llm_name=="cohere":
        llm=Cohere(
            model=model_name,
            max_tokens=700
        ) 
        history_length=200
    elif llm_name=="ai21":
        llm=AI21(
            temperature=0.7,
            model=model_name,
        )   
    elif llm_name=="nlpcloud":
        llm=NLPCloud(
            model_name=model_name,
        )
    else:
        raise Exception("Unknown LLM "+llm_name)

    print("Use model ",model_name,"from",llm_name)

    memory=ConversationSummaryBufferMemory(llm=llm, max_token_limit=history_length,human_prefix="QUESTION",ai_prefix="ANSWER", memory_key="history", input_key="question")
    chain = load_qa_with_sources_chain(
        llm,
        memory=memory, 
        prompt=prompt, 
        verbose=True,
    )

    return chain


    
def extractQuestionData(question,wordSalad):
    shortQuestion=Summary.summarizeMarkdown(question,min_length=100,max_length=1024,withCodeBlocks=False)

    context=Summary.summarizeText(wordSalad,min_length=20,max_length=32)
    keywords=[]
    keywords.extend(Summary.getKeywords(shortQuestion,2))
    keywords.extend(Summary.getKeywords(Summary.summarizeText(wordSalad,min_length=10,max_length=20),3))

    return [question,shortQuestion,context,keywords,wordSalad]


def queryChain(chain,question):
    wordSalad=""
    for h in chain.memory.buffer: wordSalad+=h+" "
    wordSalad+=" "+question    
    [question,shortQuestion,context,keywords,wordSalad]=utils.enqueue(lambda :extractQuestionData(question,wordSalad))
    affineDocs=utils.enqueue(lambda :getAffineDocs(question,context,keywords,shortQuestion,wordSalad))
    print("Found ",len(affineDocs), " affine docs")       
    print("Q: ", shortQuestion)
    output=chain({"input_documents": affineDocs, "question": shortQuestion}, return_only_outputs=True)    
    print("A :",output)
    return output


sessions={}
langchain.llm_cache = SmartCache(CONFIG)#SQLiteCache(database_path=CONFIG["CACHE_PATH"]+"/langchain.db")

def clearSessions():
    while True:
        time.sleep(60*5)
        for session in sessions:
            if sessions[session]["timeout"] < time.time():
                del sessions[session]
threading.Thread(target=clearSessions).start()

def createSessionSecret():
    hex_chars = string.hexdigits
    timeHash=hashlib.sha256(str(time.time()).encode("utf-8")).hexdigest()[:12]
    return ''.join(secrets.choice(hex_chars) for i in range(64))+timeHash


app = Flask(__name__)    

@app.route("/langs")
def langs():
    return json.dumps(Translator.getLangs())

@app.route("/session",methods = ['POST'])
def session():
    body=request.get_json()
    lang=body["lang"] if "lang" in body  else "en"
    if lang=="auto":
        lang="en"


    if not "sessionSecret" in body or body["sessionSecret"].strip()=="":
        sessionSecret=createSessionSecret()
    else:
        sessionSecret=body["sessionSecret"]

    if sessionSecret not in sessions:
        sessions[sessionSecret]={
            "chain": createChain(),
            "timeout": time.time()+60*30
        }
    else:
        sessions[sessionSecret]["timeout"]=time.time()+60*30
    welcomeText=""
    welcomeText+=Translator.translate("en", lang,"Hi there! I'm an AI assistant for the open source game engine jMonkeyEngine. I can help you with questions related to the jMonkeyEngine source code, documentation, and other related topics.")
    welcomeText+="<br><br>"
    welcomeText+="<footer><span class=\"material-symbols-outlined\">tips_and_updates</span><span>"+Translator.translate("en", lang,"This chat bot is intended to provide helpful information, but accuracy is not guaranteed.")+"</span></footer>"

       
    return json.dumps( {
        "sessionSecret": sessionSecret,
        "helloText":Translator.translate("en",lang,"Who are you?"),
        "welcomeText":welcomeText
    })

@app.route("/query",methods = ['POST'])
def query():
    try:
        body=request.get_json()
        question=rewrite(body["question"])
        lang=body["lang"] if "lang" in body  else "en"
        
        if lang == "auto":
            lang=Translator.detect(question)

        if lang!="en":
            question=Translator.translate(lang,"en",question)

        if len(question)==0:
            raise Exception("Question is empty")
        
        sessionSecret=body["sessionSecret"]   
        
        if sessionSecret not in sessions:
            return json.dumps({"error": "Session expired"})
            
        chain=sessions[sessionSecret]["chain"]

        output=queryChain(chain,question)
       
        if lang!="en":
            output["output_text"]=Translator.translate("en",lang,output["output_text"])

        #print(chain.memory.buffer)
        return json.dumps(output)
    except Exception as e:
        print(e)
        print(traceback.format_exc())
        errorStr=str(e)
        errorStr=rewriteError(errorStr)
        return json.dumps({"error": errorStr})


@app.route('/<path:filename>')
def serveFrontend(filename):
    return send_from_directory('frontend/', filename)

@app.route('/')
def serveIndex():
    return send_from_directory('frontend/', "index.html")

@app.route('/docs', methods=['POST'])
def docs():
    body=request.get_json()
    question=body["question"]
    maxFragmentsToReturn=int(body.get("maxFragmentsToReturn",3))
    maxFragmentsToSelect=int(body.get("maxFragmentsToReturn",6))
    wordSalad=body.get("context","")+" "+question
    [question,shortQuestion,context,keywords,wordSalad]=utils.enqueue(lambda : extractQuestionData(question,wordSalad))
    affineDocs=utils.enqueue(lambda : getAffineDocs(
        question,context,keywords,shortQuestion,wordSalad,
        maxFragmentsToReturn=maxFragmentsToReturn,
        maxFragmentsToSelect=maxFragmentsToSelect
    ))
    plainDocs=[
        {
            "content":doc.page_content,
            "metadata":doc.metadata
        } for doc in affineDocs
    ]
    return json.dumps(plainDocs)


serve(app, host="0.0.0.0", port=8080, connection_limit=1000)

langchain.chains.conversation.memory.ConversationSummaryBufferMemory langchain.llms.OpenAI langchain.llms.AI21 langchain.llms.Cohere langchain.chains.qa_with_sources.load_qa_with_sources_chain langchain.llms.NLPCloud langchain.prompts.PromptTemplate
[((5785, 5803), 'SmartCache.SmartCache', 'SmartCache', (['CONFIG'], {}), '(CONFIG)\n', (5795, 5803), False, 'from SmartCache import SmartCache\n'), ((6330, 6345), 'flask.Flask', 'Flask', (['__name__'], {}), '(__name__)\n', (6335, 6345), False, 'from flask import Flask, send_from_directory\n'), ((9830, 9890), 'waitress.serve', 'serve', (['app'], {'host': '"""0.0.0.0"""', 'port': '(8080)', 'connection_limit': '(1000)'}), "(app, host='0.0.0.0', port=8080, connection_limit=1000)\n", (9835, 9890), False, 'from waitress import serve\n'), ((1263, 1275), 'json.load', 'json.load', (['f'], {}), '(f)\n', (1272, 1275), False, 'import json\n'), ((1280, 1310), 'embeddings.EmbeddingsManager.init', 'EmbeddingsManager.init', (['CONFIG'], {}), '(CONFIG)\n', (1302, 1310), False, 'from embeddings import EmbeddingsManager\n'), ((1315, 1335), 'Summary.Summary.init', 'Summary.init', (['CONFIG'], {}), '(CONFIG)\n', (1327, 1335), False, 'from Summary import Summary\n'), ((1591, 1614), 'translator.Translator.init', 'Translator.init', (['CONFIG'], {}), '(CONFIG)\n', (1606, 1614), False, 'from translator import Translator\n'), ((2557, 2659), 're.sub', 're.sub', (['"""\\\\b(app|applet|game|application)\\\\b"""', '"""simple application"""', 'question'], {'flags': 're.IGNORECASE'}), "('\\\\b(app|applet|game|application)\\\\b', 'simple application',\n    question, flags=re.IGNORECASE)\n", (2563, 2659), False, 'import re\n'), ((3341, 3432), 'langchain.prompts.PromptTemplate', 'PromptTemplate', ([], {'input_variables': "['history', 'question', 'summaries']", 'template': 'template'}), "(input_variables=['history', 'question', 'summaries'],\n    template=template)\n", (3355, 3432), False, 'from langchain.prompts import PromptTemplate\n'), ((4371, 4540), 'langchain.chains.conversation.memory.ConversationSummaryBufferMemory', 'ConversationSummaryBufferMemory', ([], {'llm': 'llm', 'max_token_limit': 'history_length', 'human_prefix': '"""QUESTION"""', 'ai_prefix': '"""ANSWER"""', 'memory_key': '"""history"""', 'input_key': '"""question"""'}), "(llm=llm, max_token_limit=history_length,\n    human_prefix='QUESTION', ai_prefix='ANSWER', memory_key='history',\n    input_key='question')\n", (4402, 4540), False, 'from langchain.chains.conversation.memory import ConversationSummaryBufferMemory, ConversationBufferMemory, ConversationBufferWindowMemory\n'), ((4543, 4618), 'langchain.chains.qa_with_sources.load_qa_with_sources_chain', 'load_qa_with_sources_chain', (['llm'], {'memory': 'memory', 'prompt': 'prompt', 'verbose': '(True)'}), '(llm, memory=memory, prompt=prompt, verbose=True)\n', (4569, 4618), False, 'from langchain.chains.qa_with_sources import load_qa_with_sources_chain\n'), ((4748, 4842), 'Summary.Summary.summarizeMarkdown', 'Summary.summarizeMarkdown', (['question'], {'min_length': '(100)', 'max_length': '(1024)', 'withCodeBlocks': '(False)'}), '(question, min_length=100, max_length=1024,\n    withCodeBlocks=False)\n', (4773, 4842), False, 'from Summary import Summary\n'), ((4849, 4911), 'Summary.Summary.summarizeText', 'Summary.summarizeText', (['wordSalad'], {'min_length': '(20)', 'max_length': '(32)'}), '(wordSalad, min_length=20, max_length=32)\n', (4870, 4911), False, 'from Summary import Summary\n'), ((6497, 6515), 'flask.request.get_json', 'request.get_json', ([], {}), '()\n', (6513, 6515), False, 'from flask import request\n'), ((7046, 7280), 'translator.Translator.translate', 'Translator.translate', (['"""en"""', 'lang', '"""Hi there! I\'m an AI assistant for the open source game engine jMonkeyEngine. I can help you with questions related to the jMonkeyEngine source code, documentation, and other related topics."""'], {}), '(\'en\', lang,\n    "Hi there! I\'m an AI assistant for the open source game engine jMonkeyEngine. I can help you with questions related to the jMonkeyEngine source code, documentation, and other related topics."\n    )\n', (7066, 7280), False, 'from translator import Translator\n'), ((8890, 8932), 'flask.send_from_directory', 'send_from_directory', (['"""frontend/"""', 'filename'], {}), "('frontend/', filename)\n", (8909, 8932), False, 'from flask import Flask, send_from_directory\n'), ((8979, 9025), 'flask.send_from_directory', 'send_from_directory', (['"""frontend/"""', '"""index.html"""'], {}), "('frontend/', 'index.html')\n", (8998, 9025), False, 'from flask import Flask, send_from_directory\n'), ((9086, 9104), 'flask.request.get_json', 'request.get_json', ([], {}), '()\n', (9102, 9104), False, 'from flask import request\n'), ((9806, 9827), 'json.dumps', 'json.dumps', (['plainDocs'], {}), '(plainDocs)\n', (9816, 9827), False, 'import json\n'), ((1359, 1382), 'query.embeddingsquery.EmbeddingsQuery', 'EmbeddingsQuery', (['CONFIG'], {}), '(CONFIG)\n', (1374, 1382), False, 'from query.embeddingsquery import EmbeddingsQuery\n'), ((1392, 1545), 'query.discoursequery.DiscourseQuery', 'DiscourseQuery', (['CONFIG', "CONFIG['JME_HUB_URL']"], {'searchFilter': "CONFIG['JME_HUB_SEARCH_FILTER']", 'knowledgeCutoff': "CONFIG['JME_HUB_KNOWLEDGE_CUTOFF']"}), "(CONFIG, CONFIG['JME_HUB_URL'], searchFilter=CONFIG[\n    'JME_HUB_SEARCH_FILTER'], knowledgeCutoff=CONFIG[\n    'JME_HUB_KNOWLEDGE_CUTOFF'])\n", (1406, 1545), False, 'from query.discoursequery import DiscourseQuery\n'), ((1943, 1954), 'time.time', 'time.time', ([], {}), '()\n', (1952, 1954), False, 'import time\n'), ((3165, 3194), 'os.path.exists', 'os.path.exists', (['template_path'], {}), '(template_path)\n', (3179, 3194), False, 'import os\n'), ((3764, 3841), 'langchain.llms.OpenAI', 'OpenAI', ([], {'temperature': 'temperature', 'model_name': 'model_name', 'max_tokens': 'max_tokens'}), '(temperature=temperature, model_name=model_name, max_tokens=max_tokens)\n', (3770, 3841), False, 'from langchain.llms import OpenAI\n'), ((4946, 4983), 'Summary.Summary.getKeywords', 'Summary.getKeywords', (['shortQuestion', '(2)'], {}), '(shortQuestion, 2)\n', (4965, 4983), False, 'from Summary import Summary\n'), ((5914, 5932), 'time.sleep', 'time.sleep', (['(60 * 5)'], {}), '(60 * 5)\n', (5924, 5932), False, 'import time\n'), ((6061, 6099), 'threading.Thread', 'threading.Thread', ([], {'target': 'clearSessions'}), '(target=clearSessions)\n', (6077, 6099), False, 'import threading\n'), ((6407, 6428), 'translator.Translator.getLangs', 'Translator.getLangs', ([], {}), '()\n', (6426, 6428), False, 'from translator import Translator\n'), ((7801, 7819), 'flask.request.get_json', 'request.get_json', ([], {}), '()\n', (7817, 7819), False, 'from flask import request\n'), ((8605, 8623), 'json.dumps', 'json.dumps', (['output'], {}), '(output)\n', (8615, 8623), False, 'import json\n'), ((3930, 3970), 'langchain.llms.Cohere', 'Cohere', ([], {'model': 'model_name', 'max_tokens': '(700)'}), '(model=model_name, max_tokens=700)\n', (3936, 3970), False, 'from langchain.llms import Cohere\n'), ((5024, 5086), 'Summary.Summary.summarizeText', 'Summary.summarizeText', (['wordSalad'], {'min_length': '(10)', 'max_length': '(20)'}), '(wordSalad, min_length=10, max_length=20)\n', (5045, 5086), False, 'from Summary import Summary\n'), ((6992, 7003), 'time.time', 'time.time', ([], {}), '()\n', (7001, 7003), False, 'import time\n'), ((7398, 7532), 'translator.Translator.translate', 'Translator.translate', (['"""en"""', 'lang', '"""This chat bot is intended to provide helpful information, but accuracy is not guaranteed."""'], {}), "('en', lang,\n    'This chat bot is intended to provide helpful information, but accuracy is not guaranteed.'\n    )\n", (7418, 7532), False, 'from translator import Translator\n'), ((7636, 7684), 'translator.Translator.translate', 'Translator.translate', (['"""en"""', 'lang', '"""Who are you?"""'], {}), "('en', lang, 'Who are you?')\n", (7656, 7684), False, 'from translator import Translator\n'), ((7971, 7998), 'translator.Translator.detect', 'Translator.detect', (['question'], {}), '(question)\n', (7988, 7998), False, 'from translator import Translator\n'), ((8044, 8086), 'translator.Translator.translate', 'Translator.translate', (['lang', '"""en"""', 'question'], {}), "(lang, 'en', question)\n", (8064, 8086), False, 'from translator import Translator\n'), ((8290, 8330), 'json.dumps', 'json.dumps', (["{'error': 'Session expired'}"], {}), "({'error': 'Session expired'})\n", (8300, 8330), False, 'import json\n'), ((8499, 8554), 'translator.Translator.translate', 'Translator.translate', (['"""en"""', 'lang', "output['output_text']"], {}), "('en', lang, output['output_text'])\n", (8519, 8554), False, 'from translator import Translator\n'), ((8785, 8816), 'json.dumps', 'json.dumps', (["{'error': errorStr}"], {}), "({'error': errorStr})\n", (8795, 8816), False, 'import json\n'), ((2241, 2252), 'time.time', 'time.time', ([], {}), '()\n', (2250, 2252), False, 'import time\n'), ((4072, 4111), 'langchain.llms.AI21', 'AI21', ([], {'temperature': '(0.7)', 'model': 'model_name'}), '(temperature=0.7, model=model_name)\n', (4076, 4111), False, 'from langchain.llms import AI21\n'), ((6010, 6021), 'time.time', 'time.time', ([], {}), '()\n', (6019, 6021), False, 'import time\n'), ((6267, 6292), 'secrets.choice', 'secrets.choice', (['hex_chars'], {}), '(hex_chars)\n', (6281, 6292), False, 'import secrets\n'), ((6911, 6922), 'time.time', 'time.time', ([], {}), '()\n', (6920, 6922), False, 'import time\n'), ((8682, 8704), 'traceback.format_exc', 'traceback.format_exc', ([], {}), '()\n', (8702, 8704), False, 'import traceback\n'), ((4193, 4224), 'langchain.llms.NLPCloud', 'NLPCloud', ([], {'model_name': 'model_name'}), '(model_name=model_name)\n', (4201, 4224), False, 'from langchain.llms import NLPCloud\n'), ((6201, 6212), 'time.time', 'time.time', ([], {}), '()\n', (6210, 6212), False, 'import time\n')]