import streamlit as st
import os
import requests
import pandas as pd
import boto3
import AWSHandler
import pinecone
from sentence_transformers import SentenceTransformer

aws_access_key = os.getenv("aws_access_key")
aws_secret_key = os.getenv("aws_secret_key")
pinecone_api_key = os.getenv("pinecone_api_key")
pinecone_environment = os.getenv("pinecone_environment")

s3 = boto3.client('s3', aws_access_key_id=aws_access_key, aws_secret_access_key=aws_secret_key)

model = SentenceTransformer('all-mpnet-base-v2')

st.set_page_config(layout="wide")



def search_index(query):
    pinecone.init(api_key=pinecone_api_key, environment=pinecone_environment)
    index = pinecone.Index("scotus")
    vector = model.encode([query]).tolist()
    print(index)
    print(index.query)
    print(index.query(vector=vector,top_k=10,include_values=False,include_metadata=True).to_dict())
    return index.query(vector=vector,top_k=10,include_values=False,include_metadata=True).to_dict()['matches']
    

st.markdown('''
<style>
[data-testid="stMarkdownContainer"] ul{
    list-style-position: inside;
}
</style>
''', unsafe_allow_html=True)


search = st.text_input('Search', placeholder='First Amendment cannot protect libel')
if search != "":
    res = search_index(query=search)
    for m in res:
        text =  m['metadata']['Text']
        if 'Name' in m['metadata']:
            name = m['metadata']['Name']
        else:
            name = ""
        st.markdown("- " + text + ' ' + '(_' + name + '_)')