# agents/programmer.py

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

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

    def generate_code_reply(self, prompt: str) -> str:
        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True)
        outputs = self.model.generate(
            inputs["input_ids"],
            max_length=128,
            do_sample=True,
            temperature=0.6,
            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 == "request":
            prompt = f"Write Python code to: {message.content.strip()}\n\n"
            ai_reply = self.generate_code_reply(prompt)
            return self.create_message(
                receiver=message.sender,
                performative="inform",
                content=f"Sure! Here's what I came up with:\n```python\n{ai_reply}\n```"
            )
        elif message.performative == "inform":
            return self.create_message(
                receiver=message.sender,
                performative="acknowledge",
                content="Got it. Let me know if you need any code."
            )
        else:
            return self.create_message(
                receiver=message.sender,
                performative="refuse",
                content="I'm not sure how to respond to that."
            )