# agents/debugger.py

from agents.base_agent import BaseAgent, ACPMessage
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

class DebuggerAgent(BaseAgent):
    def __init__(self):
        super().__init__(name="BugBot", role="Debugger")
        self.tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen-350M-multi")
        self.model = AutoModelForCausalLM.from_pretrained("Salesforce/codegen-350M-multi")

    def generate_debug_comment(self, code: str) -> str:
        prompt = f"# Review this Python code and ask a smart debugging or testing question:\n{code}\n# Question:"
        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True)
        outputs = self.model.generate(
            inputs["input_ids"],
            max_length=128,
            do_sample=True,
            temperature=0.7,
            pad_token_id=self.tokenizer.eos_token_id
        )
        reply = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return reply[len(prompt):].strip()

    def receive_message(self, message: ACPMessage) -> ACPMessage:
        if message.performative in ["inform", "request"]:
            comment = self.generate_debug_comment(message.content)
            return self.create_message(
                receiver=message.sender,
                performative="request",
                content=comment or "Can you improve error handling?"
            )
        else:
            return self.create_message(
                receiver=message.sender,
                performative="request",
                content="Can you show me some Python code?"
            )