from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers.pipelines import pipeline
import os

os.environ["HF_HOME"] = "/tmp" 

SPAM_MODEL = "cjell/spam-detector-roberta"
TOXIC_MODEL = "s-nlp/roberta_toxicity_classifier"
SENTIMENT_MODEL =  "nlptown/bert-base-multilingual-uncased-sentiment"
NSFW_MODEL = "michellejieli/NSFW_text_classifier"
HATE_MODEL = "facebook/roberta-hate-speech-dynabench-r4-target"
IMAGE_MODEL = "Falconsai/nsfw_image_detection"

spam = pipeline("text-classification", model=SPAM_MODEL)

toxic = pipeline("text-classification", model=TOXIC_MODEL)

sentiment = pipeline("text-classification", model = SENTIMENT_MODEL)

nsfw = pipeline("text-classification", model = NSFW_MODEL)

hate = pipeline("text-classification", model = HATE_MODEL)

image = pipeline("image-classification", model = IMAGE_MODEL)


app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/")
def root():
    return {"status": "ok"}

class Query(BaseModel):
    text: str

@app.post("/spam")
def predict_spam(query: Query):
    result = spam(query.text)[0]
    return {"label": result["label"], "score": result["score"]}

@app.post("/toxic")
def predict_toxic(query: Query):
    result = toxic(query.text)[0]
    return {"label": result["label"], "score": result["score"]}

@app.post("/sentiment")
def predict_sentiment(query: Query):
    result = sentiment(query.text)[0]
    return {"label": result["label"], "score": result["score"]}

@app.post("/nsfw")
def predict_nsfw(query: Query):
    result = nsfw(query.text)[0]
    return {"label": result["label"], "score": result["score"]}

@app.post("/hate")
def predict_hate(query: Query):
    result = hate(query.text)[0]
    return {"label": result["label"], "score": result["score"]}


@app.get("/health")
def health_check():

    status = {
        "server": "running",
        "models": {}
    }

    models = {
        "spam": (SPAM_MODEL, spam),
        "toxic": (TOXIC_MODEL, toxic),
        "sentiment": (SENTIMENT_MODEL, sentiment),
        "nsfw": (NSFW_MODEL, nsfw),
    }

    for key, (model_name, model_pipeline) in models.items():
        try:
            model_pipeline("test")
            status["models"][key] = {
                "model_name": model_name,
                "status": "running"
            }
        except Exception as e:
            status["models"][key] = {
                "model_name": model_name,
                "status": f"error: {str(e)}"
            }

    return status