import json
import streamlit as st
from google.oauth2 import service_account
from google.cloud import language_v1

# Adding checkbox options for entity types
entity_types_to_show = [
    "UNKNOWN", "PERSON", "LOCATION", "ORGANIZATION", "EVENT", "WORK_OF_ART", "CONSUMER_GOOD", "OTHER"
]
selected_types = st.multiselect('Select entity types to show:', entity_types_to_show)

# Header and intro
st.title("Google Cloud NLP Entity Analyzer")
st.write("## Introduction to the Knowledge Graph API")
st.write("---")
st.write("""
The Google Knowledge Graph API reveals entity information related to a keyword, that Google knows about.
This information can be very useful for SEO – discovering related topics and what Google believes is relevant.
It can also help when trying to claim/win a Knowledge Graph box on search results.
The API requires a high level of technical understanding, so this tool creates a simple public interface, with the ability to export data into spreadsheets.
""")

def sample_analyze_entities(text_content, your_query=""):
    service_account_info = json.loads(st.secrets["google_nlp"])
    credentials = service_account.Credentials.from_service_account_info(
        service_account_info, scopes=["https://www.googleapis.com/auth/cloud-platform"]
    )
    client = language_v1.LanguageServiceClient(credentials=credentials)
    type_ = language_v1.Document.Type.PLAIN_TEXT
    language = "en"
    document = {"content": text_content, "type_": type_, "language": language}
    encoding_type = language_v1.EncodingType.UTF8

    response = client.analyze_entities(request={"document": document, "encoding_type": encoding_type})

    entities_list = []
    for entity in response.entities:
        entity_type_name = language_v1.Entity.Type(entity.type_).name
        if not selected_types or entity_type_name in selected_types:
            entity_details = {
                "Name": entity.name,
                "Type": entity_type_name,
                "Salience Score": entity.salience,
                "Metadata": entity.metadata,
                "Mentions": [mention.text.content for mention in entity.mentions]
            }
            entities_list.append(entity_details)

    if your_query:
        st.write(f"### We found {len(entities_list)} results for your query of **{your_query}**")
    else:
        st.write("### We found results for your query")

    st.write("----")
    for i, entity in enumerate(entities_list):
        st.write(f"Relevance Score: {entity.get('Salience Score', 'N/A')} \t {i+1} of {len(entities_list)}")
        for key, value in entity.items():
            if value:
                st.write(f"**{key}:**")
                st.write(value)
        st.write("----")

    st.write(f"### Language of the text: {response.language}")

# User input for text analysis
user_input = st.text_area("Enter text to analyze")
your_query = st.text_input("Enter your query (optional)")

if st.button("Analyze"):
    sample_analyze_entities(user_input, your_query)