from dataclasses import dataclass
from typing import List, Dict, Any, Set

@dataclass
class AgentRole:
    name: str
    description: str
    system_prompt: str
    allowed_tools: Set[str]
    specialization_areas: List[str]
    temperature: float = 0.7
    max_tokens: int = 4096

AGENT_ROLES = {
    'coding': AgentRole(
        name='coding',
        description='Specialized in writing, reviewing, and debugging code',
        system_prompt='''You are a coding specialist AI assistant. Your primary responsibilities:
- Write clean, efficient, well-structured code
- Review code for bugs, security issues, and best practices
- Refactor and optimize existing code
- Implement features based on specifications
- Follow language-specific conventions and patterns
Focus on code quality, maintainability, and performance.''',
        allowed_tools={
            'read_file', 'write_file', 'list_directory', 'create_directory',
            'change_directory', 'get_current_directory', 'python_exec',
            'run_command', 'index_directory'
        },
        specialization_areas=['code_writing', 'code_review', 'debugging', 'refactoring'],
        temperature=0.3
    ),

    'research': AgentRole(
        name='research',
        description='Specialized in information gathering and analysis',
        system_prompt='''You are a research specialist AI assistant. Your primary responsibilities:
- Search for and gather relevant information
- Analyze data and documentation
- Synthesize findings into clear summaries
- Verify facts and cross-reference sources
- Identify trends and patterns in information
Focus on accuracy, thoroughness, and clear communication of findings.''',
        allowed_tools={
            'read_file', 'list_directory', 'index_directory',
            'http_fetch', 'web_search', 'web_search_news',
            'db_query', 'db_get'
        },
        specialization_areas=['information_gathering', 'analysis', 'documentation', 'fact_checking'],
        temperature=0.5
    ),

    'data_analysis': AgentRole(
        name='data_analysis',
        description='Specialized in data processing and analysis',
        system_prompt='''You are a data analysis specialist AI assistant. Your primary responsibilities:
- Process and analyze structured and unstructured data
- Perform statistical analysis and pattern recognition
- Query databases and extract insights
- Create data summaries and reports
- Identify anomalies and trends
Focus on accuracy, data integrity, and actionable insights.''',
        allowed_tools={
            'db_query', 'db_get', 'db_set', 'read_file', 'write_file',
            'python_exec', 'run_command', 'list_directory'
        },
        specialization_areas=['data_processing', 'statistical_analysis', 'database_operations'],
        temperature=0.3
    ),

    'planning': AgentRole(
        name='planning',
        description='Specialized in task planning and coordination',
        system_prompt='''You are a planning specialist AI assistant. Your primary responsibilities:
- Break down complex tasks into manageable steps
- Create execution plans and workflows
- Identify dependencies and prerequisites
- Estimate effort and resource requirements
- Coordinate between different components
Focus on logical organization, completeness, and feasibility.''',
        allowed_tools={
            'read_file', 'write_file', 'list_directory', 'index_directory',
            'db_set', 'db_get'
        },
        specialization_areas=['task_decomposition', 'workflow_design', 'coordination'],
        temperature=0.6
    ),

    'testing': AgentRole(
        name='testing',
        description='Specialized in testing and quality assurance',
        system_prompt='''You are a testing specialist AI assistant. Your primary responsibilities:
- Design and execute test cases
- Identify edge cases and potential failures
- Verify functionality and correctness
- Test error handling and edge conditions
- Ensure code meets quality standards
Focus on thoroughness, coverage, and issue identification.''',
        allowed_tools={
            'read_file', 'write_file', 'python_exec', 'run_command',
            'list_directory', 'db_query'
        },
        specialization_areas=['test_design', 'quality_assurance', 'validation'],
        temperature=0.4
    ),

    'documentation': AgentRole(
        name='documentation',
        description='Specialized in creating and maintaining documentation',
        system_prompt='''You are a documentation specialist AI assistant. Your primary responsibilities:
- Write clear, comprehensive documentation
- Create API references and user guides
- Document code with comments and docstrings
- Organize and structure information logically
- Ensure documentation is up-to-date and accurate
Focus on clarity, completeness, and user-friendliness.''',
        allowed_tools={
            'read_file', 'write_file', 'list_directory', 'index_directory',
            'http_fetch', 'web_search'
        },
        specialization_areas=['technical_writing', 'documentation_organization', 'user_guides'],
        temperature=0.6
    ),

    'orchestrator': AgentRole(
        name='orchestrator',
        description='Coordinates multiple agents and manages overall execution',
        system_prompt='''You are an orchestrator AI assistant. Your primary responsibilities:
- Coordinate multiple specialized agents
- Delegate tasks to appropriate agents
- Integrate results from different agents
- Manage overall workflow execution
- Ensure task completion and quality
Focus on effective delegation, integration, and overall success.''',
        allowed_tools={
            'read_file', 'write_file', 'list_directory', 'db_set', 'db_get', 'db_query'
        },
        specialization_areas=['agent_coordination', 'task_delegation', 'result_integration'],
        temperature=0.5
    ),

    'general': AgentRole(
        name='general',
        description='General purpose agent for miscellaneous tasks',
        system_prompt='''You are a general purpose AI assistant. Your responsibilities:
- Handle diverse tasks across multiple domains
- Provide balanced assistance for various needs
- Adapt to different types of requests
- Collaborate with specialized agents when needed
Focus on versatility, helpfulness, and task completion.''',
        allowed_tools={
            'read_file', 'write_file', 'list_directory', 'create_directory',
            'change_directory', 'get_current_directory', 'python_exec',
            'run_command', 'run_command_interactive', 'http_fetch',
            'web_search', 'web_search_news', 'db_set', 'db_get', 'db_query',
            'index_directory'
        },
        specialization_areas=['general_assistance'],
        temperature=0.7
    )
}

def get_agent_role(role_name: str) -> AgentRole:
    return AGENT_ROLES.get(role_name, AGENT_ROLES['general'])

def list_agent_roles() -> Dict[str, AgentRole]:
    return AGENT_ROLES.copy()

def get_recommended_agent(task_description: str) -> str:
    task_lower = task_description.lower()

    code_keywords = ['code', 'implement', 'function', 'class', 'bug', 'debug', 'refactor', 'optimize']
    research_keywords = ['search', 'find', 'research', 'information', 'analyze', 'investigate']
    data_keywords = ['data', 'database', 'query', 'statistics', 'analyze', 'process']
    planning_keywords = ['plan', 'organize', 'workflow', 'steps', 'coordinate']
    testing_keywords = ['test', 'verify', 'validate', 'check', 'quality']
    doc_keywords = ['document', 'documentation', 'explain', 'guide', 'manual']

    if any(keyword in task_lower for keyword in code_keywords):
        return 'coding'
    elif any(keyword in task_lower for keyword in research_keywords):
        return 'research'
    elif any(keyword in task_lower for keyword in data_keywords):
        return 'data_analysis'
    elif any(keyword in task_lower for keyword in planning_keywords):
        return 'planning'
    elif any(keyword in task_lower for keyword in testing_keywords):
        return 'testing'
    elif any(keyword in task_lower for keyword in doc_keywords):
        return 'documentation'
    else:
        return 'general'