from _chroma import ChromaDB
from _openai import LLM_OpenAI as LLM
import streamlit as st
from print_color import *

def get_docs(user_input):
    docs = chroma.query('mala_interrogations', user_input, n_results=5)
    return docs

def generate_prompt(user_input, docs):
    texts = [text for text in docs['documents'][0]]
    metas = [{'person': meta['person'], 'date': meta['date']} for meta in docs['metadatas'][0]]
    combined_data = list(zip(texts, metas))
    string = ''
    for text, meta in combined_data:
        # Do something with text and meta
        # For example, print them
        string += f'\n\nFrån förhör med {meta["person"]} {meta["date"]}:'.upper()
        string += f'\n{text}\n\n'

    prompt = f'''Svara på frågan: {user_input}\n
    Använd endast informationen nedan:\n
    {string}\n
    Skriv utförligt på svenska och var noga med detaljer som namn, plats och datum. 
    Får gärna med information från alla fem förhör om det är relevant.\n
    {user_input}'''
    answer = st.session_state.llm.generate(prompt)
    return prompt

st.set_page_config(
    page_title="Malå",
)

# Should not be reseted every run.
if "llm" not in st.session_state:
    st.session_state.llm = LLM(chat=True, system_prompt='Du är assistent åt en journalist som går igenom förhör i en förundersökning. Svara bara utifrån den information du får. Svara alltid på svenska!')

# Should be reset every run.
llm_checker = LLM(chat=True)
chroma = ChromaDB()


# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display chat messages from history on app rerun
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Accept user input
if user_input := st.chat_input("Fråga något om förhören."):
    print_blue(user_input)
    if len(st.session_state.messages) > 1:
        history = ''
        for message in st.session_state.messages:
            history += f"{message['role']}: {message['content']}\n'"
        prompt = f'En användare har ställt frågan "{user_input}" och här är chatthistoriken mellan användaren och en assistent:\n{history}\n\nVerkar "{user_input}" vara en uppföljningfråga eller en fristående fråga? Svara ENDAST med "uppföljning" eller "fristående".'
        chat_completion = llm_checker.generate(prompt, stream=False)
        answer = chat_completion.choices[0].message.content

        print_red(answer)
        if 'uppföljning' in answer:
            prompt=f'Använd historiken till att omformulera "{user_input}" till en helt fristående fråga. Frågan ska användas för att hitta information i förhören.'
            chat_completion = llm_checker.generate(prompt, stream=False)
            question2chroma = chat_completion.choices[0].message.content
        if 'fristående' in answer:
            question2chroma = user_input

        if 'None' in answer:
            question2chroma = user_input

        print_yellow(question2chroma)

    else:
        question2chroma = user_input
    # Add user message to chat history
    st.session_state.messages.append({"role": "user", "content": user_input})
    # Display user message in chat message container
    with st.chat_message("user"):
        st.markdown(user_input)

    # Display assistant response in chat message container
    with st.chat_message("assistant"):
        docs = get_docs(question2chroma)
        prompt = generate_prompt(user_input, docs)
        stream = st.session_state.llm.generate(prompt)
        response = st.write_stream(stream)
        st.session_state.llm.messages.append({'role': 'assistant', 'content': response})
    st.session_state.messages.append({"role": "assistant", "content": response})
    print()