import streamlit as st
from openai import OpenAI
import time
import os
import re
import pandas as pd
import PyPDF2
from datetime import datetime
from pydub import AudioSegment
from docx import Document
from io import BytesIO

st.set_page_config(page_title="Schlager ContractAi")

st.title("Schlager ContractAi")
st.caption("Chat with your contract or manage meeting minutes")

# Sidebar for API Key input
with st.sidebar:
    OPENAI_API_KEY = st.text_input("Enter your C2 Group of Technologies Access Key", type="password")

# Tabs for Contract and Minutes
tab1, tab2 = st.tabs(["Contract", "Minutes"])

SUPPORTED_AUDIO_FORMATS = (".mp3", ".wav", ".m4a")
SUPPORTED_TEXT_FORMATS = (".txt", ".docx", ".csv", ".xlsx", ".pdf")

with tab1:
    st.subheader("Contract Chat")
    
    if OPENAI_API_KEY:
        client = OpenAI(api_key=OPENAI_API_KEY)
    else:
        st.error("Please enter your C2 Group of Technologies Access Key to continue.")
        st.stop()

    ASSISTANT_ID = "asst_rd9h8PfYuOmHbkvOF3RTmVfn"
    
    if "messages" not in st.session_state:
        st.session_state["messages"] = []
    
    if st.button("Clear Chat"):
        st.session_state["messages"] = []
        st.rerun()
    
    for message in st.session_state.messages:
        role, content = message["role"], message["content"]
        st.chat_message(role).write(content)

    if prompt := st.chat_input():
        st.session_state.messages.append({"role": "user", "content": prompt})
        st.chat_message("user").write(prompt)

        try:
            thread = client.beta.threads.create()
            thread_id = thread.id
            client.beta.threads.messages.create(
                thread_id=thread_id,
                role="user",
                content=prompt
            )
            
            run = client.beta.threads.runs.create(
                thread_id=thread_id,
                assistant_id=ASSISTANT_ID
            )

            while True:
                run_status = client.beta.threads.runs.retrieve(thread_id=thread_id, run_id=run.id)
                if run_status.status == "completed":
                    break
                time.sleep(1)

            messages = client.beta.threads.messages.list(thread_id=thread_id)
            assistant_message = messages.data[0].content[0].text.value
            st.chat_message("assistant").write(assistant_message)
            st.session_state.messages.append({"role": "assistant", "content": assistant_message})
        except Exception as e:
            st.error(f"Error: {str(e)}")

with tab2:
    st.subheader("Minutes")
    
    if "generated_minutes" not in st.session_state:
        st.session_state["generated_minutes"] = ""

    uploaded_files = st.file_uploader("Upload meeting minutes (PDF/DOCX/Audio)", 
                                      type=["pdf", "docx", "mp3", "wav", "m4a"], 
                                      accept_multiple_files=True)
    
    if uploaded_files:
        st.write("### Uploaded Files:")
        for uploaded_file in uploaded_files:
            st.write(f"- {uploaded_file.name}")
        
        combined_text = ""
        for uploaded_file in uploaded_files:
            if uploaded_file.name.lower().endswith(SUPPORTED_AUDIO_FORMATS):
                audio = AudioSegment.from_file(uploaded_file)
                temp_audio_path = "temp_audio.mp3"
                audio.export(temp_audio_path, format="mp3")
                with open(temp_audio_path, "rb") as audio_file:
                    transcription = client.audio.transcriptions.create(
                        model="whisper-1",
                        file=audio_file
                    )
                combined_text += transcription.text + "\n"
                os.remove(temp_audio_path)
            else:
                if uploaded_file.name.endswith(".docx"):
                    doc = Document(uploaded_file)
                    combined_text += "\n".join([para.text for para in doc.paragraphs])
                elif uploaded_file.name.endswith(".pdf"):
                    pdf_reader = PyPDF2.PdfReader(uploaded_file)
                    combined_text += "\n".join([page.extract_text() for page in pdf_reader.pages if page.extract_text()])
        
        if combined_text:
            st.write("### Transcribed and Extracted Text:")
            st.text_area("Meeting Transcript", combined_text, height=300)
            
            if st.button("Generate Meeting Minutes"):
                response = client.chat.completions.create(
                    model="gpt-4-turbo",
                    messages=[
                        {"role": "system", "content": "You are an AI assistant that generates professional meeting minutes."},
                        {"role": "user", "content": f"Summarize the following into structured meeting minutes:\n{combined_text}"}
                    ]
                )
                st.session_state["generated_minutes"] = response.choices[0].message.content
                st.write("### Meeting Minutes:")
                st.text_area("Generated Minutes", st.session_state["generated_minutes"], height=300)
            
            if st.session_state["generated_minutes"]:
                doc = Document()
                doc.add_paragraph(st.session_state["generated_minutes"])
                
                docx_io = BytesIO()
                doc.save(docx_io)
                docx_io.seek(0)
                
                st.download_button(label="Download Meeting Minutes", 
                                   data=docx_io, 
                                   file_name="Meeting_Minutes.docx", 
                                   mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document")
                
                if st.button("Contract Analysis"):
                    analysis_prompt = "Analyse this data and provide any insights, risks and relevant data based on the contract that we need to be aware of.\n" + st.session_state["generated_minutes"]
                    thread = client.beta.threads.create()
                    thread_id = thread.id
                    client.beta.threads.messages.create(
                        thread_id=thread_id,
                        role="user",
                        content=analysis_prompt
                    )
                    
                    run = client.beta.threads.runs.create(
                        thread_id=thread_id,
                        assistant_id=ASSISTANT_ID
                    )
                    
                    while True:
                        run_status = client.beta.threads.runs.retrieve(thread_id=thread_id, run_id=run.id)
                        if run_status.status == "completed":
                            break
                        time.sleep(1)
                    
                    messages = client.beta.threads.messages.list(thread_id=thread_id)
                    analysis_output = messages.data[0].content[0].text.value
                    
                    st.write("### Contract Analysis Output:")
                    st.text_area("Contract Analysis", analysis_output, height=300)