import os
os.environ['HF_HOME'] = '/tmp'
import time
import streamlit as st
import pandas as pd
import io
import plotly.express as px
import zipfile
import json
from cryptography.fernet import Fernet
from streamlit_extras.stylable_container import stylable_container
from typing import Optional
from gliner import GLiNER
from comet_ml import Experiment

# --- CSS Styling for the App ---
st.markdown(
    """
    <style>
    /* Main app background and text color */
    .stApp {
        background-color: #E8F5E9; /* A very light green */
        color: #1B5E20; /* Dark green for the text */
    }

    /* Sidebar background color */
    .css-1d36184 {
        background-color: #A5D6A7; /* A medium light green */
        secondary-background-color: #A5D6A7;
    }

    /* Expander background color and header */
    .streamlit-expanderContent, .streamlit-expanderHeader {
        background-color: #E8F5E9;
    }

    /* Text Area background and text color */
    .stTextArea textarea {
        background-color: #81C784; /* A slightly darker medium green */
        color: #1B5E20; /* Dark green for text */
    }

    /* Button background and text color */
    .stButton > button {
        background-color: #81C784;
        color: #1B5E20;
    }

    /* Warning box background and text color */
    .stAlert.st-warning {
        background-color: #66BB6A; /* A medium-dark green for the warning box */
        color: #1B5E20;
    }

    /* Success box background and text color */
    .stAlert.st-success {
        background-color: #66BB6A; /* A medium-dark green for the success box */
        color: #1B5E20;
    }
    </style>
    """,
    unsafe_allow_html=True
)

# --- Page Configuration and UI Elements ---
st.set_page_config(layout="wide", page_title="Named Entity Recognition App")
st.subheader("PiiGuard", divider="violet")
st.link_button("by nlpblogs", "https://nlpblogs.com", type="tertiary")
expander = st.expander("**Important notes**")
expander.write("""**Named Entities:** This PiiGuard web app predicts fifty-one (51) labels: "person", "organization", "social_media_handle", "username", "insurance_company", "phone_number", "email", "email_address", "mobile_phone_number", "landline_phone_number", "fax_number", "credit_card_number", "credit_card_expiration_date", "credit_card_brand", "cvv", "cvc", "bank_account_number", "iban", "transaction_number", "cpf", "cnpj",  "passport_number", "passport_expiration_date", "driver's_license_number", "tax_identification_number", "identity_card_number", "national_id_number", "identity_document_number", "birth_certificate_number", "social_security_number", "health_insurance_id_number", "health_insurance_number", "national_health_insurance_number", "student_id_number", "registration_number", "insurance_number", "serial_number", "visa_number", "reservation_number", "train_ticket_number", "medication", "medical_condition", "blood_type",  "date_of_birth", "address", "ip_address", "postal_code",  "flight_number", "license_plate_number", "vehicle_registration_number", "digital_signature"

Results are presented in easy-to-read tables, visualized in an interactive tree map, pie chart and bar chart, and are available for download along with a Glossary of tags.

**How to Use:** Type or paste your text into the text area below, then press Ctrl + Enter. Click the 'Results' button to extract and tag entities in your text data.

**Usage Limits:** You can request results unlimited times for one (1) month.

**Supported Languages:** English

**Technical issues:** If your connection times out, please refresh the page or reopen the app's URL. 

For any errors or inquiries, please contact us at info@nlpblogs.com""")

with st.sidebar:
    st.write("Use the following code to embed the PiiGuard web app on your website. Feel free to adjust the width and height values to fit your page.")
    code = '''
    <iframe
	src="https://aiecosystem-piiguard.hf.space"
	frameborder="0"
	width="850"
	height="450"
    ></iframe>

    '''
    st.code(code, language="html")
    st.text("")
    st.text("")
    st.divider()
    st.subheader("🚀 Ready to build your own AI Web App?", divider="violet")
    st.link_button("AI Web App Builder", "https://nlpblogs.com/build-your-named-entity-recognition-app/", type="primary")

# --- Comet ML Setup ---
COMET_API_KEY = os.environ.get("COMET_API_KEY")
COMET_WORKSPACE = os.environ.get("COMET_WORKSPACE")
COMET_PROJECT_NAME = os.environ.get("COMET_PROJECT_NAME")
comet_initialized = bool(COMET_API_KEY and COMET_WORKSPACE and COMET_PROJECT_NAME)
if not comet_initialized:
    st.warning("Comet ML not initialized. Check environment variables.")

# --- Label Definitions ---
labels = [
    "person", "organization", "social_media_handle", "username", "insurance_company",
    "phone_number", "email", "email_address", "mobile_phone_number", "landline_phone_number", "fax_number",
    "credit_card_number", "credit_card_expiration_date", "credit_card_brand", "cvv", "cvc",
    "bank_account_number", "iban", "transaction_number", "cpf", "cnpj",
    "passport_number", "passport_expiration_date", "driver's_license_number", "tax_identification_number", "identity_card_number",
    "national_id_number", "identity_document_number", "birth_certificate_number", "social_security_number",
    "health_insurance_id_number", "health_insurance_number", "national_health_insurance_number",
    "student_id_number", "registration_number", "insurance_number", "serial_number",
    "visa_number", "reservation_number", "train_ticket_number",
    "medication", "medical_condition", "blood_type", "date_of_birth",
    "address", "ip_address", "postal_code",
    "flight_number", "license_plate_number", "vehicle_registration_number", "digital_signature"
]

# Corrected mapping dictionary
category_mapping = {
    "People_and_Groups": [
        "person", "organization", "social_media_handle", "username", "insurance_company"
    ],
    "Contact_Information": [
        "phone_number", "email", "email_address", "mobile_phone_number", "landline_phone_number", "fax_number"
    ],
    "Financial_and_Transactions": [
        "credit_card_number", "credit_card_expiration_date", "credit_card_brand", "cvv", "cvc",
        "bank_account_number", "iban", "transaction_number", "cpf", "cnpj"
    ],
    "Identification_and_Documents": [
        "passport_number", "passport_expiration_date", "driver's_license_number", "tax_identification_number",
        "identity_card_number", "national_id_number", "identity_document_number", "birth_certificate_number",
        "social_security_number", "health_insurance_id_number", "health_insurance_number",
        "national_health_insurance_number", "student_id_number", "registration_number",
        "insurance_number", "serial_number", "visa_number", "reservation_number", "train_ticket_number"
    ],
    "Health_and_Personal": [
        "medication", "medical_condition", "blood_type", "date_of_birth"
    ],
    "Locations_and_Addresses": [
        "address", "ip_address", "postal_code"
    ],
    "Transportation_and_Logistics": [
        "flight_number", "license_plate_number", "vehicle_registration_number"
    ],
    "Digital_and_Security": [
        "digital_signature"
    ]
}

# --- Model Loading ---
@st.cache_resource
def load_ner_model():
    """Loads the GLiNER model and caches it."""
    try:
        return GLiNER.from_pretrained("knowledgator/gliner-multitask-v1.0", nested_ner=True, num_gen_sequences=2, gen_constraints=labels)
    except Exception as e:
        st.error(f"Failed to load NER model. Please check your internet connection or model availability: {e}")
        st.stop()
model = load_ner_model()

# Flatten the mapping to a single dictionary
reverse_category_mapping = {label: category for category, label_list in category_mapping.items() for label in label_list}

# --- Session State Initialization ---
if 'show_results' not in st.session_state:
    st.session_state.show_results = False
if 'last_text' not in st.session_state:
    st.session_state.last_text = ""
if 'results_df' not in st.session_state:
    st.session_state.results_df = pd.DataFrame()
if 'elapsed_time' not in st.session_state:
    st.session_state.elapsed_time = 0.0

# --- Text Input and Clear Button ---
word_limit = 200
text = st.text_area(f"Type or paste your text below (max {word_limit} words), and then press Ctrl + Enter", height=250, key='my_text_area')
word_count = len(text.split())
st.markdown(f"**Word count:** {word_count}/{word_limit}")

def clear_text():
    """Clears the text area and hides results."""
    st.session_state['my_text_area'] = ""
    st.session_state.show_results = False
    st.session_state.last_text = ""
    st.session_state.results_df = pd.DataFrame()
    st.session_state.elapsed_time = 0.0
st.button("Clear text", on_click=clear_text)

# --- Results Section ---
if st.button("Results"):
    if not text.strip():
        st.warning("Please enter some text to extract entities.")
        st.session_state.show_results = False
    elif word_count > word_limit:
        st.warning(f"Your text exceeds the {word_limit} word limit. Please shorten it to continue.")
        st.session_state.show_results = False
    else:
        if text != st.session_state.last_text:
            st.session_state.show_results = True
            st.session_state.last_text = text
            start_time = time.time()
            with st.spinner("Extracting entities...", show_time=True):
                entities = model.predict_entities(text, labels)
                df = pd.DataFrame(entities)
                st.session_state.results_df = df
                
                # Move the Comet ML logging and termination here
                if not df.empty:
                    df['category'] = df['label'].map(reverse_category_mapping)
                    if comet_initialized:
                        experiment = Experiment(
                            api_key=COMET_API_KEY,
                            workspace=COMET_WORKSPACE,
                            project_name=COMET_PROJECT_NAME,
                        )
                        experiment.log_parameter("input_text", text)
                        experiment.log_table("predicted_entities", df)
                        experiment.end()
            
            end_time = time.time()
            st.session_state.elapsed_time = end_time - start_time
        else:
            st.session_state.show_results = True

# Display results if the state variable is True
if st.session_state.show_results:
    df = st.session_state.results_df
    
    if not df.empty:
        df['category'] = df['label'].map(reverse_category_mapping)
        st.subheader("Grouped Entities by Category", divider="violet")
        
        # Create tabs for each category
        category_names = sorted(list(category_mapping.keys()))
        category_tabs = st.tabs(category_names)
        
        for i, category_name in enumerate(category_names):
            with category_tabs[i]:
                df_category_filtered = df[df['category'] == category_name]
                if not df_category_filtered.empty:
                    st.dataframe(df_category_filtered.drop(columns=['category']), use_container_width=True)
                else:
                    st.info(f"No entities found for the '{category_name}' category.")
        
        with st.expander("See Glossary of tags"):
            st.write('''
            - **text**: ['entity extracted from your text data']
            - **score**: ['accuracy score; how accurately a tag has been assigned to a given entity']
            - **label**: ['label (tag) assigned to a given extracted entity']
            - **start**: ['index of the start of the corresponding entity']
            - **end**: ['index of the end of the corresponding entity']
            ''')
        st.divider()
        
        # Tree map
        st.subheader("Tree map", divider="violet")
        fig_treemap = px.treemap(df, path=[px.Constant("all"), 'category', 'label', 'text'], values='score', color='category')
        fig_treemap.update_layout(margin=dict(t=50, l=25, r=25, b=25), paper_bgcolor='#E8F5E9', plot_bgcolor='#E8F5E9')
        st.plotly_chart(fig_treemap)
        
        # Pie and Bar charts
        grouped_counts = df['category'].value_counts().reset_index()
        grouped_counts.columns = ['category', 'count']
        col1, col2 = st.columns(2)
        
        with col1:
            st.subheader("Pie chart", divider="violet")
            fig_pie = px.pie(grouped_counts, values='count', names='category', hover_data=['count'], labels={'count': 'count'}, title='Percentage of predicted categories')
            fig_pie.update_traces(textposition='inside', textinfo='percent+label')
            fig_pie.update_layout(
                paper_bgcolor='#E8F5E9',
                plot_bgcolor='#E8F5E9'
            )
            st.plotly_chart(fig_pie)
            
        with col2:
            st.subheader("Bar chart", divider="violet")
            fig_bar = px.bar(grouped_counts, x="count", y="category", color="category", text_auto=True, title='Occurrences of predicted categories')
            fig_bar.update_layout(
                paper_bgcolor='#E8F5E9',
                plot_bgcolor='#E8F5E9'
            )
            st.plotly_chart(fig_bar)
            
        # Most Frequent Entities
        st.subheader("Most Frequent Entities", divider="gray")
        word_counts = df['text'].value_counts().reset_index()
        word_counts.columns = ['Entity', 'Count']
        repeating_entities = word_counts[word_counts['Count'] > 1]
        
        if not repeating_entities.empty:
            st.dataframe(repeating_entities, use_container_width=True)
            fig_repeating_bar = px.bar(repeating_entities, x='Entity', y='Count', color='Entity')
            fig_repeating_bar.update_layout(xaxis={'categoryorder': 'total descending'},
                                            paper_bgcolor='#E8F5E9',
                                            plot_bgcolor='#E8F5E9')
            st.plotly_chart(fig_repeating_bar)
        else:
            st.warning("No entities were found that occur more than once.")
        
        # Download Section
        st.divider()
        dfa = pd.DataFrame(
            data={
                'Column Name': ['text', 'label', 'score', 'start', 'end'],
                'Description': [
                    'entity extracted from your text data',
                    'label (tag) assigned to a given extracted entity',
                    'accuracy score; how accurately a tag has been assigned to a given entity',
                    'index of the start of the corresponding entity',
                    'index of the end of the corresponding entity',
                ]
            }
        )
        buf = io.BytesIO()
        with zipfile.ZipFile(buf, "w") as myzip:
            myzip.writestr("Summary of the results.csv", df.to_csv(index=False))
            myzip.writestr("Glossary of tags.csv", dfa.to_csv(index=False))
        
        with stylable_container(
            key="download_button",
            css_styles="""button { background-color: #81C784; border: 1px solid black; padding: 5px; color: #1B5E20; }""",
        ):
            st.download_button(
                label="Download results and glossary (zip)",
                data=buf.getvalue(),
                file_name="nlpblogs_results.zip",
                mime="application/zip",
            )
            
    else: # If df is empty
        st.warning("No entities were found in the provided text.")
        
    st.text("")
    st.text("")
    st.info(f"Results processed in **{st.session_state.elapsed_time:.2f} seconds**.")