from crewai import CrewBase, Agent, Task, Crew
from crewai_tools import MCPServerAdapter
from mcp import StdioServerParameters
import os
from typing import Dict, Any


class InfraReviewCrew(CrewBase):
    """Infrastructure Review Crew for conducting infrastructure reviews."""
    
    agents_config = "config/agents.yaml"
    tasks_config = "config/tasks.yaml"

    def __init__(self):
        super().__init__()
        # Configure Hadolint MCP server connection
        self.hadolint_server_params = StdioServerParameters(
            command="python",
            args=["/home/armistace/dev/pr_reviewer/mcp_servers/hadolint_mcp.py"],
            env=os.environ
        )
        # Configure Checkov MCP server connection
        self.checkov_server_params = StdioServerParameters(
            command="python",
            args=["/home/armistace/dev/pr_reviewer/mcp_servers/checkov_mcp.py"],
            env=os.environ
        )

    @Agent
    def infra_reviewer(self) -> Agent:
        """DevOps/Platform Engineer agent for infrastructure review."""
        return Agent(
            config=self.agents_config["infra_reviewer"],
            tools=[],  # Tools will be added via MCP adapter in the task
            verbose=True
        )

    @Task
    def infra_review_task(self) -> Task:
        """Task for conducting infrastructure review."""
        return Task(
            config=self.tasks_config["infra_review_task"],
        )

    @Crew
    def crew(self) -> Crew:
        """Create the Infrastructure Review crew."""
        # Create MCP server adapters for Hadolint and Checkov
        hadolint_adapter = MCPServerAdapter(self.hadolint_server_params)
        checkov_adapter = MCPServerAdapter(self.checkov_server_params)
        
        # Combine tools from both adapters
        all_tools = []
        if hasattr(hadolint_adapter, 'tools'):
            all_tools.extend(hadolint_adapter.tools)
        if hasattr(checkov_adapter, 'tools'):
            all_tools.extend(checkov_adapter.tools)
        
        return Crew(
            agents=[self.infra_reviewer()],
            tasks=[self.infra_review_task()],
            process="sequential",
            verbose=True,
            tools=all_tools,
        )