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

class TranscriptionAgent:
    def __init__(self):
        self.llm = ChatOpenAI(model="gpt-4o", temperature=0)
        
    def __call__(self, state: CallState) -> CallState:
        """Converts audio to text using whisper."""
        file_path = state["file_path"]
        
        # Load whisper model - using 'base' for faster processing
        model = whisper.load_model("base")
        result = model.transcribe(file_path)
        state["content"] = result["text"]
        
        # Clean the transcript
        prompt = PromptTemplate.from_template(
            "Clean the following text of any profanity and fix basic grammatical errors. Return only the clean text:\n{text}"
        )
        chain = prompt | self.llm
        clean_text = chain.invoke({"text": state["content"]}).content
        state["clean_content"] = clean_text
        
        return state