from fastapi import FastAPI
from pydantic import BaseModel
import joblib
import nltk
from nltk.corpus import stopwords
from nltk.stem import PorterStemmer
import re

nltk.download('stopwords')

app = FastAPI()

# Load the model pipeline
pipeline = joblib.load('spam_classifier_pipeline.joblib')

class EmailRequest(BaseModel):
    subject: str
    body: str

def preprocess_text(text):
    text = text.lower()
    text = re.sub(r'[^a-zA-Z\s]', '', text)
    words = text.split()
    stop_words = set(stopwords.words('english'))
    words = [word for word in words if word not in stop_words]
    stemmer = PorterStemmer()
    words = [stemmer.stem(word) for word in words]
    return ' '.join(words)

@app.post("/predict")
async def predict(email: EmailRequest):
    processed_text = preprocess_text(email.subject + ' ' + email.body)
    prediction = pipeline.predict([processed_text])[0]
    return {'prediction': ['ham', 'not_spam', 'spam'][prediction]}

@app.get("/")
async def root():
    return {"message": "Spam Classification API"}