import streamlit as st
import requests
import json
import re

API_URL = "https://jj78yliqd7.execute-api.eu-central-1.amazonaws.com/dev/chatbot"


def get_response(prompt, history):
    params = {"prompt": prompt}
    response = requests.get(API_URL, params=params)
    if response.status_code == 200:
        response_body = response.json().get('body')
        answer = response_body.get('answer')
        return answer
    else:
        return "Sorry, something went wrong. Please try again."


# Regular expression to extract JSON object from text
def extract_json(text):
    json_data = re.search(r'\{.*\}', text, re.DOTALL)
    if json_data:
        return json_data.group(0)
    else:
        return None

# Streamlit App
st.sidebar.title("Wissensguru - IDTA")

# Session state initialization
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# Displaying conversation history
for message in st.session_state.chat_history:
    if isinstance(message, AIMessage):
        with st.chat_message("AI", expanded=False):
            st.markdown(message.content)
           
    elif isinstance(message, HumanMessage):
        with st.chat_message("Human", expanded=False):
            st.markdown(message.content)
           

# User input
user_query = st.chat_input("Type your message here...")
if user_query is not None and user_query != "":
    # Display user's message
    with st.chat_message("Human"):
        st.markdown(user_query)
    
    # Generate bot's response
    response = get_response(user_query, st.session_state.chat_history)
    
    # Extract JSON from bot's response
    extracted_json = extract_json(response)
    
    if extracted_json:
        # Parse JSON string into a dictionary
        data_dict = json.loads(extracted_json)
        
        # Display JSON in Streamlit
        st.json(data_dict)
    else:
        # Display bot's response if JSON extraction fails
        with st.chat_message("AI"):
            st.markdown(response)
    
    # Append messages to chat history
    st.session_state.chat_history.append(HumanMessage(content=user_query))
    st.session_state.chat_history.append(AIMessage(content=response))