import os
import streamlit as st
from pinecone import Pinecone
from sentence_transformers import SentenceTransformer

# Title of the Streamlit App
st.title("Pinecone Query Search on 'pubmed-splade' Index")

# Initialize Pinecone globally
index = None

# Function to initialize Pinecone
def initialize_pinecone():
    api_key = os.getenv('PINECONE_API_KEY')  # Get Pinecone API key from environment variable
    if api_key:
        # Initialize Pinecone client using the new class instance method
        return Pinecone(api_key=api_key)
    else:
        st.error("Pinecone API key not found! Please set the PINECONE_API_KEY environment variable.")
        return None

# Function to connect to the 'pubmed-splade' index
def connect_to_index(pc):
    index_name = 'pubmed-splade'  # Hardcoded index name
    if index_name in pc.list_indexes().names():
        st.info(f"Successfully connected to index '{index_name}'")
        return pc.Index(index_name)
    else:
        st.error(f"Index '{index_name}' not found!")
        return None

# Function to encode query using sentence transformers model
def encode_query(model, query_text):
    return model.encode(query_text).tolist()

# Initialize Pinecone
pc = initialize_pinecone()

# If Pinecone initialized successfully, proceed with index management
if pc:
    index = connect_to_index(pc)

    # Load model for query encoding
    model = SentenceTransformer('msmarco-bert-base-dot-v5')

    # Query input
    query_text = st.text_input("Enter a Query to Search", "Can clinicians use the PHQ-9 to assess depression?")
    
    # Button to encode query and search the Pinecone index
    if st.button("Search Query"):
        if query_text and index:
            dense_vector = encode_query(model, query_text)

            # Search the index
            results = index.query(
                vector=dense_vector,
                top_k=5,
                include_metadata=True
            )
            
            st.write("### Search Results:")
            if results.matches:
                for match in results.matches:
                    score = match.score
                    context = match.metadata.get("context", "No context available")
                    
                    # Display score and context in a formatted way
                    st.markdown(f"**Score**: `{score}`")
                    st.markdown(f"**Context**: {context}")
                    st.markdown("---")  # Divider for each result
            else:
                st.warning("No results found for this query.")
        else:
            st.error("Please enter a query and ensure the index is initialized.")