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

# Set page title and layout
st.set_page_config(page_title="CSV News Classifier", layout="wide")

# Load the fine-tuned Hugging Face model
@st.cache_resource
def load_model():
    return pipeline("text-classification", model="TAgroup5/daily-mirror-news-classifier")

classifier = load_model()

# Custom CSS for Colors and Styling
st.markdown(
    """
    <style>
        body { background-color: #f8f9fa; }
        .stApp { background-color: #ffffff; padding: 20px; border-radius: 10px; box-shadow: 0px 0px 10px rgba(0,0,0,0.1); }
        
        .custom-box {
            background-color: #f0f0f0; 
            padding: 15px;
            border-radius: 10px;
            border: 1px solid #ccc;
            margin-bottom: 15px;
        }
        h1 { color: #ff5733; text-align: center; }  
        h2, h3 { color: #007bff; }  
    </style>
    """,
    unsafe_allow_html=True
)

# Page Title
st.title("📊 News Classification from CSV")

# File Uploader
st.markdown('<div class="custom-box"><h3>📂 Upload a CSV file</h3></div>', unsafe_allow_html=True)
uploaded_file = st.file_uploader("", type=["csv"])

if uploaded_file is not None:
    # Read CSV
    df = pd.read_csv(uploaded_file)

    # Show Preview
    st.markdown("### 🔍 **Preview of Uploaded File**", unsafe_allow_html=True)
    st.dataframe(df.head())

    # Assuming the column with news articles is named "news_text"
    if "news_text" not in df.columns:
        st.error("❌ The uploaded CSV must contain a 'news_text' column.")
    else:
        # Perform classification
        st.markdown("### 🏷️ **Classifying News Articles...**")
        with st.spinner("Processing..."):
            df["class"] = df["news_text"].apply(lambda text: classifier(text)[0]["label"])

        # Show Preview of Results
        st.markdown("### 📌 **Preview of Classified Data**", unsafe_allow_html=True)
        st.dataframe(df.head())

        # Download Button for Classified CSV
        csv = df.to_csv(index=False).encode("utf-8")
        st.markdown("### 📥 **Download Classified CSV**", unsafe_allow_html=True)
        st.download_button(
            label="⬇️ **Download Classified CSV**",
            data=csv,
            file_name="classified_news.csv",
            mime="text/csv",
            help="Click to download the classified news file"
        )

# Footer
st.markdown("---")
st.markdown('<p style="text-align:center; font-size:14px; color:#6c757d;">👨‍💻 Developed by <b>Ridmi Navodya</b> | Powered by Streamlit 🚀</p>', unsafe_allow_html=True)