from openai import OpenAI
from gradio.chat_interface import ChatInterface
import os
from dotenv import load_dotenv
from gradio.themes.soft import Soft
from gradio.themes import colors
from gradio.components import Textbox
from gradio.components import Chatbot
from RAG_class import RAG_1177
load_dotenv()

textbox = Textbox(placeholder="Skriv ditt fråga här...",scale=4)
chatbot = Chatbot(placeholder="Varsågod och ställ din fråga nedanför 😊",scale=3, height=250, show_copy_button=True, label="1177 chatbot")

new_primary_color = colors.red
new_secondary_color = colors.red

my_custom_theme = Soft(
    primary_hue=new_primary_color,
    secondary_hue=new_secondary_color
)

# System prompt
system_prompt = """
Välkommen till 1177 AI-chatbot!

Jag är här för att hjälpa dig med sjukvårdsfrågor genom att ge svar baserade på dokumentation från 1177.se. Nedan ser du de olika typerna av frågor jag kan svara på:
   
1. **Relevanta frågor med svar från 1177:** Jag kan ge detaljerade svar och inkludera tydliga referenser när frågan kan besvaras med information från 1177.se.

2. **Relevanta följdfrågor:** Om du har uppföljningsfrågor kan jag använda tidigare svar som kontext för att ge mer information.

4. **Orelevanta frågor:** Vänligen håll dig till ämnet sjukvård och hälsa för bästa svar.

Låt mig veta hur jag kan hjälpa dig!

"""

examples = [
        "hur kan jag lindra min huvudvärk?",
        "Vad är symptomen för covid-19?",
        "Vilka läkemedel finns det för smärtlindring?",
        "Hur kan jag bäst lindra mina symtom på pollenallergi?",
        "Vad kan jag göra för att förbättra min sömn?",
        "vad kan jag göra om mitt barn får magsjuka?"  
    ]



description = """<ul>
  <li><strong>Specialiserad:</strong> Designad för att endast ge svar på sjukvårdsfrågor som man ofta besöker 1177 för.</li>
  <li><strong>Svar med referenser:</strong> Alla svar är baserade på data från <a href="https://www.1177.se/" target="_blank">1177-hemsidan</a> och innehåller källhänvisningar.</li>
  <li><strong>Tips</strong> var så specifik som möjligt. Du kan även använda de inbyggda frågorna.</li>

<p>Du är varmt välkommen med frågor eller feedback: hassan.saado94@gmail.com, <a href="https://www.linkedin.com/in/hassan-saado/">Linkedin</a></li></p>"""



def predict(message, history):
    # Initialize RAG model
    rag = RAG_1177()
    print(message)
    relevance_check = rag.relevant_question(message)

    # Prepare history in OpenAI format
    history_openai_format = []
    for human, assistant in history:
        history_openai_format.append({"role": "user", "content": human})
        history_openai_format.append({"role": "assistant", "content": assistant})

    documents, urls = rag.retrieve(message, 3)

    if not history:
        history_openai_format.append({"role": "system", "content": system_prompt})

    if relevance_check == "NEJ":
        message = "Denna fråga är helt orrelevant och håller sig inte till ämnet. Ge inga referenser"

    user_prompt = f"""Fråga: {message}

    Information från 1177.se:
    {documents}

    Referenser:
    {urls}

    Svara på frågan baserat på informationen ovan och ange källorna med nummer i slutet. 
    Om frågan är en uppföljningsfråga kan du strunta i den givna informationen. Svara inte på frågan om det inte är relaterat till hälsa och sjukvård.
    """
    history_openai_format.append({"role": "user", "content": user_prompt})

    client = OpenAI()
    response = client.chat.completions.create(
        model='gpt-3.5-turbo',
        messages=history_openai_format,
        temperature=0.4,
        stream=True
    )

    partial_message = ""
    for chunk in response:
        if chunk.choices[0].delta.content is not None:
            partial_message = partial_message + chunk.choices[0].delta.content
            yield partial_message

    history_openai_format.append({"role": "assistant", "content": partial_message})

    yield partial_message

def main():

    ChatInterface(predict, textbox=textbox, chatbot=chatbot, title="Välkommen till 1177 AI-chatbot! 🔍🤒", theme=my_custom_theme,submit_btn="Skicka",
                  retry_btn="🔄Försök igen", undo_btn="↩️ Ångra", clear_btn="🗑️ Rensa",
                   examples=examples,  cache_examples=False, description=description).launch(share=False)

if __name__ == "__main__":
    main()