import streamlit as st
import openai
import anthropic
import fitz

# Set up API clients
openai_api_key = st.secrets["OPENAI_API_KEY"]
openai_client = openai.OpenAI(api_key=openai_api_key)

anthropic_api_key = st.secrets["ANTHROPIC_API_KEY"]
anthropic_client = anthropic.Anthropic(api_key=anthropic_api_key)


# Password protection
def check_password():
    if "password_correct" not in st.session_state:
        st.session_state["password_correct"] = False

    if not st.session_state["password_correct"]:
        password = st.text_input("Enter the password", type="password")
        if st.button("Submit"):
            if password == st.secrets["PASSWORD"]:
                st.session_state["password_correct"] = True
                st.rerun()
            else:
                st.error("😕 Password incorrect")
        return False
    else:
        return True


def extract_text_from_pdf(file):
    try:
        with fitz.open(stream=file.read(), filetype="pdf") as doc:
            text = ""
            for page in doc:
                text += page.get_text()
        return text
    except Exception as e:
        st.error(f"Error parsing PDF: {str(e)}")
        return None


def main():
    st.title("Harvey Legal Research Take-Home")

    if not check_password():
        return

    # Model selector
    model = st.selectbox("Select Model", ["GPT-4o", "Claude 3.5 Sonnet"])

    # Prompt input
    prompt = st.text_area("Enter your prompt:", height=200)

    # Document upload
    uploaded_file = st.file_uploader("Upload a document", type=["txt", "pdf"])
    document_text = None
    if uploaded_file is not None:
        if uploaded_file.type == "text/plain":
            document_text = uploaded_file.getvalue().decode()
        elif uploaded_file.type == "application/pdf":
            document_text = extract_text_from_pdf(uploaded_file)
        else:
            st.error("Unsupported file type. Please upload a text or PDF file.")

    # Run button
    if st.button("Run"):
        if not prompt:
            st.warning("Please enter a prompt.")
            return

        if document_text is None:
            full_prompt = prompt
        else:
            full_prompt = f"Document content:\n\n{document_text}\n\nPrompt:\n\n{prompt}"

        if model == "GPT-4o":
            with st.spinner("Processing with GPT-4o..."):
                stream = openai_client.chat.completions.create(
                    model="gpt-4o",
                    messages=[{"role": "user", "content": full_prompt}],
                    stream=True,
                )

                full_response = ""
                response_area = st.empty()

                for chunk in stream:
                    if chunk.choices[0].delta.content is not None:
                        full_response += chunk.choices[0].delta.content
                        response_area.markdown(f"**Response:**\n\n{full_response}")

        elif model == "Claude 3.5 Sonnet":
            with st.spinner("Processing with Claude 3.5 Sonnet..."):
                with anthropic_client.messages.stream(
                    model="claude-3-5-sonnet-20240620",
                    max_tokens=8192,
                    messages=[{"role": "user", "content": full_prompt}],
                ) as stream:
                    full_response = ""
                    response_area = st.empty()

                    for text in stream.text_stream:
                        full_response += text
                        response_area.markdown(f"**Response:**\n\n{full_response}")


if __name__ == "__main__":
    main()