import os

import requests
import streamlit as st
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Read API URL from environment variable
API_URL = os.getenv("API_URL", "http://localhost:8000")


def get_similar_prompts(query: str, n: int) -> dict:
    """
    Fetches similar prompts from the API based on the user query.

    Args:
    query (str): The user query for which similar prompts are to be retrieved.
    n (int): The number of similar prompts to return.

    Returns:
    dict: A dictionary containing similar prompts, or None if there was an error.

    Raises:
    requests.RequestException: If an HTTP error occurs during the request.
    """
    try:
        response = requests.post(
            f"{API_URL}/most_similar", json={"query": query, "n": n}
        )
        response.raise_for_status()  # Raise an exception for HTTP errors
        return response.json()
    except requests.RequestException as e:
        st.error(f"Error: {e}")
        return {}


def get_color(score: float) -> str:
    """
    Determines the color based on the similarity score.

    Args:
    score (float): The similarity score of a prompt.

    Returns:
    str: The color representing the score, which could be "green", "orange", or "red".
    """
    if score >= 0.8:
        return "green"
    elif score >= 0.5:
        return "orange"
    else:
        return "red"


def main():
    """
    The main function for running the Streamlit app.
    Sets up the UI for entering queries and retrieving similar prompts.
    """
    st.title("Prompt Similarity Finder")

    query = st.text_input("Enter your query:", "")
    n = st.slider(
        "Number of similar prompts to retrieve:", min_value=1, max_value=40, value=5
    )

    if st.button("Find Similar Prompts"):
        if query:
            with st.spinner("Fetching similar prompts..."):
                result = get_similar_prompts(query, n)
                if result:
                    similar_prompts = result.get("similar_prompts", [])
                    if similar_prompts:
                        st.subheader("Similar Prompts:")
                        for item in similar_prompts:
                            score = item["score"]
                            color = get_color(score)
                            st.markdown(
                                f"<p><strong>Score:</strong> <span style='color:{color};'>{score:.2f}</span> <br> <strong>Prompt:</strong> {item['prompt']}</p>",
                                unsafe_allow_html=True,
                            )
                            st.write("---")
                    else:
                        st.write("No similar prompts found.")
        else:
            st.warning("Please enter a query.")


if __name__ == "__main__":
    main()