from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from src.agents.CallState import CallState

class ModerationAgent:
    """
    Agent that uses LLM to detect and redact obscene words, replacing them with ***.
    Updates the clean_content in the state.
    """

    def __init__(self):
        self.llm = ChatOpenAI(model="gpt-4o", temperature=0)

    def __call__(self, state: CallState) -> CallState:
        text = state.get("clean_content") or state.get("content") or ""
        if not text:
            return state

        prompt = PromptTemplate.from_template(
            "You are a moderation assistant. Review the following text. "
            "If you find any obscene words, profanity, or inappropriate language, replace them entirely with ***. "
            "Do not change any other words. Return only the redacted text.\n\n"
            "Text:\n{text}"
        )
        
        chain = prompt | self.llm
        result = chain.invoke({"text": text})
        
        # Save the redacted text back to clean_content
        state["clean_content"] = result.content
        return state