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, Technical, and Minutes
tab1, tab2, tab3 = st.tabs(["Contract", "Technical", "Minutes"])

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


def contract_chat_section(tab, assistant_id, session_key):
    with tab:
        st.subheader("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()

        if session_key not in st.session_state:
            st.session_state[session_key] = []

        if st.button("Clear Chat", key=f"clear_chat_{session_key}"):
            st.session_state[session_key] = []
            st.rerun()

        for message in st.session_state[session_key]:
            role, content = message["role"], message["content"]
            st.chat_message(role).write(content)

        if prompt := st.chat_input():
            st.session_state[session_key].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[session_key].append({"role": "assistant", "content": assistant_message})
            except Exception as e:
                st.error(f"Error: {str(e)}")

ASSISTANT_CONTRACT_ID = "asst_rd9h8PfYuOmHbkvOF3RTmVfn"
ASSISTANT_TECHNICAL_ID = "asst_technical_example"  # Replace with actual assistant ID

# Contract Chat Section
contract_chat_section(tab1, ASSISTANT_CONTRACT_ID, "contract_messages")

# Technical Chat Section
contract_chat_section(tab2, ASSISTANT_TECHNICAL_ID, "technical_messages")

with tab3:
    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", key="generate_minutes"):
                prompt = f"""
                Based on the following meeting transcript, generate structured meeting minutes in the format below:
                
                ---
                **Meeting Name:** [Enter meeting name]  
                **Location:** [Enter location]  
                **Date:** [Enter date]  
                **Time:** [Enter time]  
                **Attendees:** [List attendees]  
                ### Agenda Items
                - [Agenda Item 1]
                - [Agenda Item 2]
                - [Agenda Item 3]
                ### Action Items
                | Action Item | Owner(s) | Deadline | Status |
                |------------|---------|----------|--------|
                | [Action Item 1] | [Owner(s) 1] | [Deadline 1] | [Status 1] |
                | [Action Item 2] | [Owner(s) 2] | [Deadline 2] | [Status 2] |
                | [Action Item 3] | [Owner(s) 3] | [Deadline 3] | [Status 3] |
                ---
                **Meeting Summary:**  
                [Brief summary of key points, discussions, and decisions]
                **Transcript:**  
                {combined_text}
                """

                response = client.chat.completions.create(
                    model="gpt-4-turbo",
                    messages=[
                        {"role": "system", "content": "You are an AI assistant that generates professional meeting minutes in a structured template."},
                        {"role": "user", "content": prompt}
                    ]
                )
                st.session_state["generated_minutes"] = response.choices[0].message.content
                
                st.write("### Generated Meeting Minutes:")
                st.text_area("Meeting Minutes", st.session_state["generated_minutes"], height=400)