from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline
from glin_profanity import Filter

# Initialize your models
toxicity_pipe = pipeline("text-classification", model="textdetox/twitter-xlmr-toxicity-classifier")
detox_pipe = pipeline("text2text-generation", model="s-nlp/mt0-xl-detox-orpo")
profanity_filter = Filter({"languages": ["english", "spanish", "french"], "replace_with": "***"})

app = FastAPI(title="Multilingual Moderation API")

class TextRequest(BaseModel):
    text: str

class ModerationResponse(BaseModel):
    original_text: str
    detected_categories: dict
    is_flagged: bool
    rewritten_text: str | None

@app.post("/moderate", response_model=ModerationResponse)
async def moderate_text(request: TextRequest):
    text = request.text
    profanity_result = profanity_filter.check_profanity(text)
    toxicity_scores = toxicity_pipe(text)[0]
    
    is_flagged = profanity_result["contains_profanity"] or toxicity_scores['score'] > 0.5
    rewritten = None
    if is_flagged:
        rewritten = detox_pipe(text, max_length=100)[0]['generated_text']
    
    return ModerationResponse(
        original_text=text,
        detected_categories={"profanity": profanity_result["contains_profanity"], "toxicity_model": toxicity_scores},
        is_flagged=is_flagged,
        rewritten_text=rewritten
    )

@app.get("/health")
async def health_check():
    return {"status": "ok"}