import joblib
import re
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from fastapi import FastAPI
from pydantic import BaseModel

# Load the model and vectorizer
vectorizer = joblib.load("vectorizer.joblib")
model = joblib.load("naive_bayes_model.joblib")

app = FastAPI()

class URLInput(BaseModel):
    url: str

def preprocess_url(url):
    url = re.sub(r"http\S+", "", url)
    url = re.sub(r"\d+", "", url)
    url = re.sub(r"\W", " ", url)
    url = url.lower()
    return url

@app.post("/predict")
def predict_url(url_input: URLInput):
    processed_url = preprocess_url(url_input.url)
    vectorized_url = vectorizer.transform([processed_url])
    prediction = model.predict(vectorized_url)
    return {"prediction": prediction[0]}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)