import streamlit as st
from openai import OpenAI
import re
import os

# Initialize OpenAI client (REPLACE with your actual key from environment variables or a secure config)

# IMPORTANT: Never hardcode API keys in your code, especially for sharing!

client = OpenAI(

base_url="https://integrate.api.nvidia.com/v1", # Assuming this is your correct base URL

api_key="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" # Replace with your actual key

)


st.title("Sorting Hat")

def get_essay_answer():
    """Gets the essay answer from the user."""
    st.subheader("Answer the Question:")
    essay = st.text_area("Describe what you would do if you encountered a lost and distressed first-year student who has lost their pet toad and is terrified of being late for Herbology.", height=200) # Increased height for essay
    return essay

def get_sorted(essay):
    """Handles the sorting process using the OpenAI API."""
    if essay:  # Check if essay is not empty
        output_area = st.empty()

        messages = [{"role": "user", "content": f"""Pretend that you are the sorting hat in Harry Potter. Sort me into a house in Hogwarts based on my answer. Give me the house ONLY (one word: Gryffindor, Hufflepuff, Ravenclaw, Slytherin) on the first line. On the second line explain your reasoning for sorting me there. Here is my answer:
        {essay}
        """}]

        with st.spinner("Sorting..."):
            try:
                completion = client.chat.completions.create(
                    model="meta/llama-3.1-405b-instruct",  # Your model
                    messages=messages,
                    temperature=0.7,
                    max_tokens=200,
                    stream=True,
                )

                full_response = ""
                message_placeholder = st.empty()

                for chunk in completion:
                    if chunk.choices[0].delta and chunk.choices[0].delta.content:
                        full_response += chunk.choices[0].delta.content
                        message_placeholder.markdown(full_response, unsafe_allow_html=False)
                    elif chunk.choices[0].finish_reason == "stop":
                        break

                house_match = re.search(r"(Gryffindor|Hufflepuff|Ravenclaw|Slytherin)", full_response, re.IGNORECASE)

                if house_match:
                    house = house_match.group(0).strip()
                    try:
                        parts = full_response.split(house + '\n', 1)
                        if len(parts) > 1:
                            explanation = parts[1].strip()
                        else:
                            explanation = "No explanation provided."
                    except Exception as e:
                        explanation = f"Error extracting explanation: {e}"

                    gif_path = {
                        "Gryffindor": "https://i.pinimg.com/originals/52/46/eb/5246eba389dfc5722324b484975a012f.gif",
                        "Hufflepuff": "https://c.tenor.com/ThUF6Bm-GLgAAAAM/always-harry-potter.gif",
                        "Ravenclaw": "https://media.tenor.com/Sgm7TzrEQzAAAAAC/harry-potter-ravenclaw.gif",
                        "Slytherin": "https://media.giphy.com/media/A1xhQcogWZAIg/giphy.gif"
                    }.get(house, "default.gif")

                    output_area.image(gif_path, width=300)

                else:
                    output_area.error("The Sorting Hat's response was not in the expected format.")

            except Exception as e:
                output_area.error(f"An error occurred: {e}")

    else:
        st.warning("Please provide an answer.")


# Main app flow
essay = get_essay_answer()  # Get the essay answer

if st.button("Get Sorted!"):
    get_sorted(essay)  # Pass the essay to get_sorted()