import streamlit as st
import openai
import os
import base64
import glob
import json
import mistune
import pytz
import textwrap

from datetime import datetime
from openai import ChatCompletion
from xml.etree import ElementTree as ET
from bs4 import BeautifulSoup

openai.api_key = os.getenv('OPENAI_KEY')
st.set_page_config(
    page_title="GPT Streamlit Document Reasoner",
    layout="wide")

menu = ["txt", "htm", "md", "py"]
choice = st.sidebar.selectbox("Choose output file type to save results", menu)
choicePrefix = "Output and download file set to "
if choice == "txt":
     st.sidebar.write(choicePrefix + "Text file.")
elif choice == "htm":
     st.sidebar.write(choicePrefix + "HTML5.")
elif choice == "md":
     st.sidebar.write(choicePrefix + "Markdown.")
elif choice == "py":
     st.sidebar.write(choicePrefix + "Python Code.")

max_length = st.sidebar.slider("Max document length", min_value=1000, max_value=32000, value=2000, step=1000)

def truncate_document(document, length):
    return document[:length]

def chat_with_model(prompts):
    model = "gpt-3.5-turbo"
    conversation = [{'role': 'system', 'content': 'You are a helpful assistant.'}]
    conversation.extend([{'role': 'user', 'content': prompt} for prompt in prompts])
    response = openai.ChatCompletion.create(model=model, messages=conversation)
    return response['choices'][0]['message']['content']

def generate_filename(prompt, file_type):
    central = pytz.timezone('US/Central')
    safe_date_time = datetime.now(central).strftime("%m%d_%I%M")  
    safe_prompt = "".join(x for x in prompt if x.isalnum())[:28]
    return f"{safe_date_time}_{safe_prompt}.{file_type}"

def create_file(filename, prompt, response):
    if filename.endswith(".txt"):
        with open(filename, 'w') as file:
            file.write(f"Prompt:\n{prompt}\nResponse:\n{response}")
    elif filename.endswith(".htm"):
        with open(filename, 'w') as file:
            file.write(f"<h1>Prompt:</h1> <p>{prompt}</p> <h1>Response:</h1> <p>{response}</p>")
    elif filename.endswith(".md"):
        with open(filename, 'w') as file:
            file.write(f"# Prompt:\n{prompt}\n# Response:\n{response}")

def get_table_download_link(file_path):
    with open(file_path, 'r') as file:
        data = file.read()
    b64 = base64.b64encode(data.encode()).decode()  
    file_name = os.path.basename(file_path)
    ext = os.path.splitext(file_name)[1] 

    if ext == '.txt':
        mime_type = 'text/plain'
    elif ext == '.htm':
        mime_type = 'text/html'
    elif ext == '.md':
        mime_type = 'text/markdown'
    else:
        mime_type = 'application/octet-stream' 

    href = f'<a href="data:{mime_type};base64,{b64}" target="_blank" download="{file_name}">{file_name}</a>'
    return href

def CompressXML(xml_text):
    root = ET.fromstring(xml_text)
    for elem in list(root.iter()):
        if isinstance(elem.tag, str) and 'Comment' in elem.tag:
            elem.parent.remove(elem)
    return ET.tostring(root, encoding='unicode', method="xml")
    
def read_file_content(file,max_length):
    if file.type == "application/json":
        content = json.load(file)
        return str(content)
    elif file.type == "text/html" or file.type == "text/htm":
        content = BeautifulSoup(file, "html.parser")
        return content.text
    elif file.type == "application/xml" or file.type == "text/xml":
        tree = ET.parse(file)
        root = tree.getroot()
        xml = CompressXML(ET.tostring(root, encoding='unicode'))
        return xml
    elif file.type == "text/markdown" or file.type == "text/md":
        md = mistune.create_markdown()
        content = md(file.read().decode())
        return content
    elif file.type == "text/plain":
        return file.getvalue().decode()
    else:
        return ""

def main():
    user_prompt = st.text_area("Your question:", '', height=120)
    uploaded_file = st.file_uploader("Choose a file", type=["xml", "json", "html", "htm", "md", "txt"])

    if st.button('💬 Chat'):
        st.write('Thinking and Reasoning with your inputs...')
        file_content = ""
        
        if user_prompt:
            prompts = textwrap.wrap(user_prompt, max_length)
            for prompt in prompts:
                response = chat_with_model([prompt])
                st.write('Response:')
                st.write(response)
                filename = generate_filename(prompt, choice)
                create_file(filename, prompt, response)
                st.sidebar.markdown(get_table_download_link(filename), unsafe_allow_html=True)

        if uploaded_file is not None:
            file_content = read_file_content(uploaded_file, max_length)
            document_parts = textwrap.wrap(file_content, max_length)
            for part in document_parts:
                response = chat_with_model([part])
                st.write('Response:')
                st.write(response)
                filename = generate_filename(part, choice)
                create_file(filename, part, response)
                st.sidebar.markdown(get_table_download_link(filename), unsafe_allow_html=True)
        
        if len(file_content) > 0:
            st.text_area("File Content:", file_content, height=400)  # Display file content in a scrollable text box
    
    all_files = glob.glob("*.txt") + glob.glob("*.htm") + glob.glob("*.md")
    for file in all_files:
        col1, col2 = st.sidebar.columns([4,1]) 
        with col1:
            st.markdown(get_table_download_link(file), unsafe_allow_html=True)
        with col2:
            if st.button("🗑", key=file):
                os.remove(file)
                st.experimental_rerun()

if __name__ == "__main__":
    main()