#!/usr/bin/env python3
"""
Custom agent example - Add your own agents to the chain.

This example adds a Tester agent that validates the generated code.
"""

from agentchain import (
    AgentChain,
    Agent,
    AgentContext,
    AgentResult,
    agent,
    Event
)


# Method 1: Class-based agent with decorator
@agent(name="tester", role="Quality Assurance Engineer")
class TesterAgent(Agent):
    """Agent that tests generated code."""
    
    system_prompt = """You are a QA engineer. Review the code and report issues.
    
    Respond with JSON:
    {
        "tests_passed": true/false,
        "issues": ["list of issues if any"],
        "summary": "brief summary"
    }
    """
    
    def execute(self, context: AgentContext) -> AgentResult:
        # Get the files that were generated
        files = context.get("files", [])
        
        # Here you could:
        # - Run pytest
        # - Do static analysis
        # - Query LLM to review code
        
        self._logger.info(f"Testing {len(files)} files...")
        
        # For demo, just return success
        return AgentResult.ok({
            "tests_passed": True,
            "files_tested": len(files),
        })


# Method 2: Function-based agent with decorator
@agent(name="documenter", role="Documentation Writer")
def documenter_agent(context: AgentContext, agent: Agent) -> AgentResult:
    """Agent that generates documentation."""
    
    # Use LLM to generate docs
    prompt = f"Generate README for: {context.task}"
    # response = agent.query_llm(prompt)
    
    return AgentResult.ok({
        "documentation": "# Project\n\nGenerated documentation...",
    })


# Use the custom agents
def main():
    result = AgentChain() \
        .requirements("Build a URL shortener service") \
        .model("deepseek-coder-v2:16b") \
        .register(TesterAgent) \
        .register(documenter_agent) \
        .configure("tester", model="llama3:8b") \
        .output("./url_shortener") \
        .on(Event.TASK_COMPLETE, lambda e: print(f"  ✓ {e.data.get('task')}")) \
        .run()
    
    print(f"\nResult: {'Success' if result.success else 'Failed'}")


if __name__ == "__main__":
    main()
