main.py

import math
import streamlit as st

from utils.model import pii_masking_pipeline, valid_labels
from utils.redis import redis_client

st.set_page_config(page_title="PII Masking Tool", page_icon="🔒")

st.title("PII Masking Tool")

if 'selected_labels' not in st.session_state:
    st.session_state.selected_labels = ["name", "nric", "phone number", "address", "email"]

def select_all_labels():
    st.session_state.selected_labels = valid_labels

# Text input
text_to_mask = st.text_area("Enter text to mask PII:", height=200)

col1, col2 = st.columns([3, 1])
with col1:
    selected_labels = st.multiselect(
        "Select PII types to mask:",
        options=valid_labels,
        key='selected_labels', 
        help="Choose the types of PII you want to redact from the text"
    )

with col2:
    st.button("Select All", on_click=select_all_labels)


if st.button("Mask Text"):
    if text_to_mask:
        masked_text = pii_masking_pipeline(text_to_mask)

        # Display results
        st.subheader("Masked Result:")
        st.write(masked_text)

        redis_client.lpush("masked_texts", masked_text)

    else:
        st.warning("Please enter some text to mask.")


# Pagination logic
ITEMS_PER_PAGE = 10
total_items = redis_client.llen("masked_texts")
total_pages = math.ceil(total_items / ITEMS_PER_PAGE)

if total_items > 0:
    st.subheader("Previous Masked Texts")
    
    col1, col2, col3 = st.columns([1, 3, 1])
    with col1:
        page = st.number_input("Page", min_value=1, max_value=max(1, total_pages), value=1)
    with col2:
        st.write(f"Page {page} of {total_pages}")

    start_idx = (page - 1) * ITEMS_PER_PAGE
    end_idx = start_idx + ITEMS_PER_PAGE - 1

    texts = redis_client.lrange("masked_texts", start_idx, end_idx)
    for i, txt in enumerate(texts, start=start_idx + 1):
        display_txt = txt.decode('utf-8')[:30] + "..." if len(txt) > 30 else txt.decode('utf-8')
        with st.expander(f"#{i}. {display_txt}"):
            st.write(txt.decode('utf-8'))