import streamlit as st
import pickle
import re
import nltk
from nltk.corpus import stopwords
from nltk.stem import SnowballStemmer
from typing import Any

# --- Configuration ---
MODEL_FILE = "src/stress_detectipn.pkl"
VECTORIZER_FILE = "src/vectorizer.pkl"

# Make sure stopwords are available on HF Spaces
try:
    nltk.data.find("corpora/stopwords")
except LookupError:
    nltk.download("stopwords")

# Preprocessing tools
stemmer = SnowballStemmer("english")
stop_words_set = set(stopwords.words("english"))


# --- 1. Load Model + Vectorizer ---
@st.cache_resource
def load_artifacts():
    """Load trained model and vectorizer."""
    try:
        with open(MODEL_FILE, "rb") as model_file:
            model = pickle.load(model_file)

        with open(VECTORIZER_FILE, "rb") as vec_file:
            vectorizer = pickle.load(vec_file)

        return model, vectorizer

    except FileNotFoundError:
        st.error(
            f"Error: Required files ({MODEL_FILE} or {VECTORIZER_FILE}) were not found. "
            "Please upload them to the Space."
        )
        return None, None


# --- 2. Preprocessing Pipeline ---
def preprocess_text(text: str) -> str:
    """Clean and prepare text exactly as in training."""
    text = str(text).lower()

    # Remove brackets, URLs, HTML
    text = re.sub(r"\[.*?\]", "", text)
    text = re.sub(r"https?://\S+|www\.\S+", "", text)
    text = re.sub(r"<.*?>", "", text)

    # Remove punctuation and special characters
    text = re.sub(r"[^\w\s]", "", text)
    text = re.sub(r"[^a-zA-Z0-9]", " ", text)

    # Remove numbers, extra spaces
    text = re.sub(r"\w*\d\w*", "", text)
    text = re.sub(r"\s+", " ", text).strip()

    # Token list
    tokens = text.split()

    # Remove stopwords
    tokens = [w for w in tokens if w not in stop_words_set]

    # Stemming
    tokens = [stemmer.stem(w) for w in tokens]

    return " ".join(tokens)


# --- 3. Prediction Logic ---
def predict_stress(text: str, model: Any, vectorizer: Any) -> str:
    """Return model prediction as label ('Stress' / 'No Stress')."""
    if not text:
        return "Please enter text."

    cleaned = preprocess_text(text)
    X = vectorizer.transform([cleaned])

    # Your model already returns the final string label
    prediction = model.predict(X)[0]
    return prediction


# --- 4. Streamlit UI ---
model, vectorizer = load_artifacts()

if model and vectorizer:

    st.title("🧠 Stress Detection System")
    st.write("Enter any text and the model will classify it as **Stress** or **No Stress**.")

    user_input = st.text_area("Enter your text here:", "")

    if st.button("Analyze"):

        if not user_input.strip():
            st.warning("Please enter some text.")
        else:
            result = predict_stress(user_input, model, vectorizer)

            st.subheader("Result")
            if result == "Stress":
                st.error(f"Prediction: **{result}** 😥")
            elif result == "No Stress":
                st.success(f"Prediction: **{result}** 😊")
            else:
                st.info(f"Prediction: **{result}**")