from fastapi import FastAPI
from pydantic import BaseModel
import joblib
import pandas as pd
import re
from collections import Counter
from fastapi.middleware.cors import CORSMiddleware

# Initialize App
app = FastAPI(title="Enterprise Spam Classifier API")

# Allow Streamlit UI to talk to this API
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Load the trained model and our 3,000 vocabulary words into memory
model = joblib.load('spam_model.pkl')
features = joblib.load('feature_names.pkl')

# Define the expected JSON payload
class EmailPayload(BaseModel):
    content: str

@app.post("/predict")
def predict_spam(payload: EmailPayload):
    # 1. Clean and tokenize the raw incoming text
    text = payload.content.lower()
    words = re.findall(r'\b\w+\b', text)
    word_counts = Counter(words)
    
    # 2. Map the user's words to the exact 3,000 columns the model expects
    input_data = {feature: [word_counts.get(feature, 0)] for feature in features}
    df_input = pd.DataFrame(input_data)
    
    # 3. Run Inference
    prediction = model.predict(df_input)[0]
    probabilities = model.predict_proba(df_input)[0]
    confidence = probabilities.max()
    
    return {
        "prediction": "Spam" if prediction == 1 else "Not Spam",
        "confidence": float(confidence),
        "status": 200
    }