import streamlit as st
from transformers import pipeline
import pandas as pd


@st.cache_resource
def load_sentiment_model():
    return pipeline("text-classification", model="jinchenliuljc/ecommerce-sentiment-analysis")

@st.cache_resource
def load_ner_model():
    return pipeline("ner", model="jinchenliuljc/ecom_ner_model")


def extract_products(ner_result):
    products = []
    current_product = None
    
    for entity in ner_result:
        if entity['entity'] == 'B-HCCX':
            if current_product is not None:
                products.append(current_product)
            current_product = {
                'start': entity['start'],
                'end': entity['end'],
                'text': entity['word']
            }
        elif entity['entity'] == 'I-HCCX' and current_product is not None:
            current_product['end'] = entity['end']
            current_product['text'] += entity['word']
    
    if current_product is not None:
        products.append(current_product)
    
    return [p['text'] for p in products]

if 'records' not in st.session_state:
    st.session_state.records = []

st.title("DTC Customer Review Analysis System")
col1, col2 = st.columns(2)

with col1:
    user_input = st.text_input("Enter customer review:", key="input")
    
    if user_input:
        sentiment_classifier = load_sentiment_model()
        sentiment_result = sentiment_classifier(user_input)[0]['label']
        

        if sentiment_result == 'LABEL_1':
            st.success("Thank you for your positive feedback! :)")
        else:

            ner_pipe = load_ner_model()
            ner_result = ner_pipe(user_input)

            products = extract_products(ner_result)
            
            # Extract products
            products = extract_products(ner_result)
            
            if products:
                # Add to records
                for product in products:
                    new_record = {
                        'Product Category': product,
                        'Review Content': user_input
                    }
                    st.session_state.records.append(new_record)
with col2:
    if st.session_state.records:
        df = pd.DataFrame(st.session_state.records)
        st.dataframe(
            df,
            column_config={
                "Product Category": "Affected Product",
                "Review Content": "Related Review"
            },
            hide_index=True,
            use_container_width=True
        )
    else:
        st.info("No feedback records yet")