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 provided 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 provided key to continue.")
        st.stop()

    ASSISTANT_ID = "asst_rd9h8PfYuOmHbkvOF3RTmVfn"
    if "messages" not in st.session_state:
        st.session_state["messages"] = []

    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")
    
    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}"}
                    ]
                )
                minutes = response.choices[0].message.content
                st.write("### Meeting Minutes:")
                st.text_area("Generated Minutes", minutes, height=300)

                date_stamp = datetime.now().strftime("%Y-%m-%d")
                file_name = f"Minutes_{date_stamp}.docx"
                
                doc = Document()
                doc.add_paragraph(minutes)
                
                docx_io = BytesIO()
                doc.save(docx_io)
                docx_io.seek(0)
                
                st.download_button(label="Download Meeting Minutes", 
                                   data=docx_io, 
                                   file_name=file_name, 
                                   mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document")