import os
import streamlit as st
import time
import json
from datetime import datetime
from openai import OpenAI

# Streamlit Page Config
st.set_page_config(page_title="Schlager ContractAi", layout="wide")

# Authentication
AUTHORIZED_USERS = {
    "andrew@lortechnologies.com": "Pass.123",
    "manish@schlagergroup.com.au": "Pass.123",
    "BenB@schlagergroup.com.au": "Pass.123",
    "nick@schlagergroup.com.au": "Pass.123"
}

if "authenticated" not in st.session_state:
    st.session_state["authenticated"] = False

def login():
    email = st.session_state.get("email", "")
    password = st.session_state.get("password", "")
    if email in AUTHORIZED_USERS and AUTHORIZED_USERS[email] == password:
        st.session_state["authenticated"] = True
    else:
        st.error("Invalid email or password. Please try again.")

if not st.session_state["authenticated"]:
    st.title("Sign In")
    st.text_input("Email", key="email")
    st.text_input("Password", type="password", key="password")
    st.button("Login", on_click=login)
    st.stop()

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

# Load API Key from Environment Variable
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    st.error("Missing OpenAI API key. Please set it as an environment variable.")
    st.stop()

# Tabs
tab1, tab2 = st.tabs(["Contract", "Technical"])

# Function to save flagged responses
def save_flagged_response(user_query, ai_response):
    flagged_data = {
        "query": user_query,
        "response": ai_response,
        "timestamp": datetime.now().isoformat()
    }
    with open("flagged_responses.json", "a") as file:
        file.write(json.dumps(flagged_data) + "\n")
    st.success("Response flagged for review.")

# Contract Chat Section
def contract_chat_section(tab, assistant_id, session_key, input_key):
    with tab:
        st.subheader("Chat")
        client = OpenAI(api_key=OPENAI_API_KEY)

        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 idx, message in enumerate(st.session_state[session_key]):
            role, content = message["role"], message["content"]
            if role == "assistant":
                col1, col2 = st.columns([5, 1])
                with col1:
                    st.chat_message(role).write(content)
                with col2:
                    if st.button("🚩 Flag", key=f"flag_{session_key}_{idx}"):
                        save_flagged_response(st.session_state[session_key][idx-1]["content"], content)
            else:
                st.chat_message(role).write(content)

        if prompt := st.chat_input("Enter your message:", key=input_key):
            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)}")

# Assign Assistant IDs
ASSISTANT_CONTRACT_ID = "asst_rd9h8PfYuOmHbkvOF3RTmVfn"
ASSISTANT_TECHNICAL_ID = "asst_xizNZBCJuy4TqdjqjwkxbAki"

contract_chat_section(tab1, ASSISTANT_CONTRACT_ID, "contract_messages", "contract_input")
contract_chat_section(tab2, ASSISTANT_TECHNICAL_ID, "technical_messages", "technical_input")