import streamlit as st
import requests
import pandas as pd

# Set Streamlit configuration for wide mode and dark theme
st.set_page_config(layout="wide", page_title="2-2")

# Modified function to send file to the API and get the response, now includes system_message parameter
def upload_file(file):
    if file is not None:
        url = "https://insly.ai/api/docusmart?locale=en-US"
        headers = {
            "accept": "multipart/form-data",
            "X-API-Key": "ai",
        }
 
        data = {
            'system_message': '"""Respond in ALLCAPS"""',
        }
        files = {
            "file": (file.name, file, file.type)
        }
        response = requests.post(url, headers=headers, files=files, data=data)
        return response.json()
    return None

# Function to display response and its subelements including key value pairs
def display_response_and_table(response_json):
    if response_json:
        # Displaying the full JSON response in a collapsed textbox
        with st.expander("Full JSON Response", expanded=False):
            st.json(response_json)
        
        # Extracting the 'response' subelement
        response_data = response_json.get('response', {})
        
        # Checking if there are key value pairs and merging them with the rest of the response data
        key_value_pairs = response_data.pop('key_value_pairs', {})
        all_response_data = {**response_data, **key_value_pairs}
        
        # Converting the merged dictionary to a pandas DataFrame for display
        df = pd.DataFrame(list(all_response_data.items()), columns=['Key', 'Value'])
        
        # Displaying the DataFrame as a transposed table in Streamlit
        st.table(df)

    else:
        st.write("No response to display.")

# Create two columns for file uploads
col1, col2 = st.columns(2)

with col1:
    file1 = st.file_uploader("Upload file to Column 1", key="file1", label_visibility="collapsed")
    if file1:
        json_response1 = upload_file(file1)
        if json_response1:
            display_response_and_table(json_response1)

with col2:
    file2 = st.file_uploader("Upload file to Column 2", key="file2", label_visibility="collapsed")
    if file2:
        json_response2 = upload_file(file2)
        if json_response2:
            display_response_and_table(json_response2)