from crewai import Agent, Task, Crew
from crewai.project import CrewBase, agent, task, crew
from crewai_tools import MCPServerAdapter
from mcp import StdioServerParameters
import os
from typing import Dict, Any
from pr_reviewer.llm import get_llm


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

    def __init__(self):
        self.llm = get_llm()
        self.hadolint_server_params = StdioServerParameters(
            command="python",
            args=["/app/mcp_servers/hadolint_mcp.py"],
            env=os.environ
        )
        self.checkov_server_params = StdioServerParameters(
            command="python",
            args=["/app/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"],
            llm=self.llm,
            tools=[],
            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."""
        all_tools = []
        try:
            hadolint_adapter = MCPServerAdapter(self.hadolint_server_params)
            if hasattr(hadolint_adapter, 'tools'):
                all_tools.extend(hadolint_adapter.tools)
        except Exception as e:
            print(f"Hadolint MCP adapter not available: {e}")
        try:
            checkov_adapter = MCPServerAdapter(self.checkov_server_params)
            if hasattr(checkov_adapter, 'tools'):
                all_tools.extend(checkov_adapter.tools)
        except Exception as e:
            print(f"Checkov MCP adapter not available: {e}")
        
        return Crew(
            agents=[self.infra_reviewer()],
            tasks=[self.infra_review_task()],
            process="sequential",
            verbose=True,
            tools=all_tools,
        )