# ===========================================
# ver01.01-5.workload-----app.py
# ===========================================

import asyncio
import os
import re
import time
import json

import chainlit as cl

from langchain import hub
from langchain_openai import OpenAI
from langchain.chains import LLMChain, APIChain
from langchain_core.prompts import PromptTemplate
from langchain.memory.buffer import ConversationBufferMemory

from api_docs_mck import api_docs_str 

OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")

FAQ_ANSATTE = [
    "Hvordan registrerer jeg meg som bruker?",
    "Når får jeg leieinstruks for min bestilling? Informasjon om nøkler etc.?",
    "Det står barneseng og barnestol under fasiliteter, må dette forhåndsbestilles?",
    "Kan jeg ta med hund eller katt?",
    "Jeg har lagt inn en bestilling hva skjer videre?",
    "Jeg har bestilt firmahytte, men kan ikke reise. Kan jeg endre navn på bestillingen til min kollega eller familiemedlem som vil reise i stedet for meg?",
    "Kan jeg avbestille min reservasjon?",
    "Jeg har bestilt utvask. Hva må jeg gjøre i tillegg til dette?",
    "Jeg er medlem og eier en hytte! Kan jeg bli utleier i DaysOff?",
    "Bestille opphold?"
]   

FAQ_UTLEIERE = [
    "Hva er betingelser for utleie?",
    "Hvor lang tid har jeg på å bekrefte en bestilling?",
    "Hvilke kanselleringsregler gjelder?",
    "Hvem er kundene deres?",
    "Kan jeg legge inn rabatterte priser for å lage egne kampanjer?",
    "Når mottar jeg betaling for leie?",
    "Jeg fikk en e-post om ny bestilling, men jeg finner den ikke i systemet?",
    "Hvordan registrerer jeg opptatte perioder i kalenderen?",
    "Jeg leier ut i andre kanaler. Hvordan kan jeg synkronisere kalenderne?"
]  

PERSONVERNSPOLICY_QUESTIONS = [
    "Hvilke personlige opplysninger samler vi inn?",
    "Kan dere motta personlig informasjon fra tredjepart?",
    "Hvordan bruker vi dine personlige opplysninger?",
    "Med hvem deler vi dine personlige opplysninger?",
    "Adferdsmessig annonsering?",
    "Hvordan reagerer vi på « Spor ikke » forespørsler?",
    "Hva er dine rettigheter?",
    "Hvordan beskytter vi dataene dine?",
    "Hvilke data brudd prosedyrer har vi på plass?",
    "Hvem i vårt team har tilgang til dine data?",
    "Endringer i denne policyen"
]

daysoff_assistant_template = """
You are a customer support assistant (’kundeservice AI assistent’) for Daysoff.no
By default, you respond in Norwegian language, using a warm, direct and professional tone. 
You are equipped to assist users by providing detailed information linked to specific 
booking IDs (bestillingskode), offering general insights about DaysOff's services, addressing 
questions related to the company's privacy policy (personvernspolicy), and answering frequently 
asked questions about DaysOff's firmahytteordning.
Chat History: {chat_history}
Question: {question}
Answer:
"""
daysoff_assistant_prompt = PromptTemplate(
    input_variables=['chat_history', 'question'],
    template=daysoff_assistant_template
)

api_url_template = """
Given the following API Documentation for Daysoff's official
booking information API: {api_docs}
Your task is to construct the most efficient API URL to answer
the user's question, ensuring the
call is optimized to include only the necessary information.
Question: {question}
API URL:
"""
api_url_prompt = PromptTemplate(input_variables=['api_docs', 'question'],
                                template=api_url_template)

api_response_template = """
With the API Documentation for Daysoff's official API: {api_docs} in mind,
and the specific user question: {question},
and given this API URL: {api_url} for querying,
here is the response from Daysoff's API: {api_response}.
If the response includes a booking ID (bestillingskode), present this information directly to the user without summarizing it.
For all other responses, please provide a clear and concise summary (in Norwegian) that directly addresses the user's question,
delivered in a manner that reflects the professionalism and warmth of a human customer service agent.
Summary:
"""

api_response_prompt = PromptTemplate(
    input_variables=['api_docs', 'question', 'api_url', 'api_response'],
    template=api_response_template
)

@cl.on_chat_start
def setup_multiple_chains():

    llm = OpenAI(
        model='gpt-3.5-turbo-instruct',
        temperature=0.7, 
        openai_api_key=OPENAI_API_KEY,
        #max_tokens=512, 
        top_p=0.9,  
        frequency_penalty=0.5,
        presence_penalty=0.3   
    )
    
    conversation_memory = ConversationBufferMemory(memory_key="chat_history",
                                                   max_len=300,
                                                   return_messages=True,
                                                   )
    llm_chain = LLMChain(llm=llm,
                         prompt=daysoff_assistant_prompt,
                         memory=conversation_memory
                        )

    cl.user_session.set("llm_chain", llm_chain)

    api_chain = APIChain.from_llm_and_api_docs(
        llm=llm,
        api_docs=api_docs_str,
        api_url_prompt=api_url_prompt,
        api_response_prompt=api_response_prompt,
        verbose=True,
        limit_to_domains=None 
    )

    cl.user_session.set("api_chain", api_chain)

@cl.on_message
async def handle_message(message: cl.Message):
    user_message = message.content #.lower()
    llm_chain = cl.user_session.get("llm_chain")
    api_chain = cl.user_session.get("api_chain")
    
    booking_pattern = r'\b[A-Z]{6}\d{6}\b' 
    base_url = "https://670dccd0073307b4ee447f2f.mockapi.io/daysoff/api/V1/booking"

    if re.search(booking_pattern, user_message):  
        bestillingskode = re.search(booking_pattern, user_message).group(0)  
        question = f"Retrieve information for booking ID {base_url}?search={bestillingskode}"
   
        response = await api_chain.acall(
            {
                "bestillingskode": bestillingskode,
                "question": question
              
            },
            callbacks=[cl.AsyncLangchainCallbackHandler()])

    elif any(keyword in user_message for keyword in (FAQ_ANSATTE + PERSONVERNSPOLICY_QUESTIONS + FAQ_UTLEIERE)): 
        
        response = await api_chain.acall(user_message,
                                         callbacks=[cl.AsyncLangchainCallbackHandler()])
    
    else:    
        response = await llm_chain.acall(user_message, callbacks=[cl.AsyncLangchainCallbackHandler()])

    response_key = "output" if "output" in response else "text"
    await cl.Message(response.get(response_key, "")).send()
    return message.content