import streamlit as st
from models.vectorizer import Vectorizer
from models.prompt_search_engine import PromptSearchEngine
from models.data_reader import load_prompts_from_jsonl
from models.Query import Query, SimilarPrompt, SearchResponse, PromptVector, VectorResponse
from sentence_transformers import SentenceTransformer
import os

# Path to your prompts data (you need to upload this file to your Hugging Face space)
prompt_path = "models/prompts_data.jsonl"  # Update this to the correct path in your space

# Initialize search engine and model
prompts = load_prompts_from_jsonl(prompt_path)
search_engine = PromptSearchEngine()
search_engine.add_prompts_to_vector_database(prompts)

# Streamlit App Interface
st.title("Prompt Search Engine")
st.write("Search for similar prompts using the local search engine.")

# Input for the user's prompt
query_input = st.text_input("Enter your prompt:")

# Number of similar prompts to retrieve (k)
k = st.number_input("Number of similar prompts to retrieve:", min_value=1, max_value=10, value=3)

# Button to trigger search
if st.button("Search Prompts"):
    if query_input:
        query = Query(prompt=query_input)
        similar_prompts, distances = search_engine.most_similar(query.prompt, top_k=k)

        # Format and display search results
        response = [
            SimilarPrompt(prompt=prompt, distance=float(distance))
            for prompt, distance in zip(similar_prompts, distances)
        ]
        st.write("Search Results:")
        for result in response:
            st.write(f"Prompt: {result.prompt}, Distance: {result.distance}")
    else:
        st.error("Please enter a prompt.")

# Additional functionality for vector similarity
st.write("---")
st.write("### Vector Similarities")

if st.button("Retrieve All Vector Similarities"):
    if query_input:
        query = Query(prompt=query_input)
        query_embedding = search_engine.model.encode([query.prompt])  # Encode the prompt to a vector
        all_similarities = search_engine.cosine_similarity(query_embedding, search_engine.index)

        # Format and display vector similarities
        response = [
            PromptVector(vector=index, distance=float(distance))
            for index, distance in enumerate(all_similarities)
        ]
        st.write("Vector Similarities:")
        for result in response:
            st.write(f"Vector Index: {result.vector}, Distance: {result.distance}")
    else:
        st.error("Please enter a prompt.")



# # streamlit_app.py
# import streamlit as st
# import requests

# # Streamlit app title
# st.title("Top K Search with Vector DataBase")

# # FastAPI endpoint URL
# # url = "http://localhost:8084/search/"
# url = "https://huggingface.co/search/"

# # Input fields in Streamlit
# id = st.text_input("Enter ID:", value="1")
# prompt = st.text_input("Enter your prompt:")
# k = st.number_input("Top K results:", min_value=1, max_value=100, value=3)

# # Trigger the search when the button is clicked
# if st.button("Search"):
#     # Construct the request payload
#     payload = {
#         "id": id,
#         "prompt": prompt,
#         "k": k
#     }

#     # Make the POST request
#     response = requests.post(url, json=payload)

#     # Handle the response
#     if response.status_code == 200:
#         results = response.json()
#         st.write(results)
#     else:
#         st.error(f"Error: {response.status_code} - {response.text}")