#!/usr/bin/env python3
"""
Ralph Agent Capability Registry

Maintains a comprehensive registry of all available agents (contains-studio, custom, etc.)
and their capabilities for dynamic selection and routing.
"""

import json
import os
import re
from typing import Dict, List, Optional, Set
from dataclasses import dataclass, field
from enum import Enum
import logging

logger = logging.getLogger('ralph.registry')


class AgentCategory(Enum):
    """Categories of agents"""
    ENGINEERING = "engineering"
    DESIGN = "design"
    PRODUCT = "product"
    MARKETING = "marketing"
    PROJECT_MANAGEMENT = "project-management"
    STUDIO_OPERATIONS = "studio-operations"
    TESTING = "testing"
    BONUS = "bonus"


class TriggerType(Enum):
    """How agents can be triggered"""
    EXPLICIT = "explicit"           # User mentions agent by name
    KEYWORD = "keyword"             # Triggered by specific keywords
    CONTEXT = "context"             # Triggered by project context
    PROACTIVE = "proactive"         # Automatically triggers
    FILE_PATTERN = "file_pattern"   # Triggered by file operations


@dataclass
class AgentCapability:
    """Represents a single agent's capabilities"""
    name: str
    category: AgentCategory
    description: str
    keywords: List[str] = field(default_factory=list)
    trigger_types: List[TriggerType] = field(default_factory=list)
    file_patterns: List[str] = field(default_factory=list)
    tools: List[str] = field(default_factory=list)
    examples: List[Dict] = field(default_factory=list)
    confidence_threshold: float = 0.5
    priority: int = 5  # 1-10, higher = preferred


class AgentCapabilityRegistry:
    """
    Registry for all available agents and their capabilities

    Maintains:
    - Agent metadata and descriptions
    - Trigger keywords and patterns
    - Tool access requirements
    - Usage statistics
    - Performance metrics
    """

    def __init__(self, agents_dir: Optional[str] = None):
        """Initialize the registry"""
        self.agents_dir = agents_dir or os.path.expanduser('~/.claude/agents')
        self.agents: Dict[str, AgentCapability] = {}
        self.keyword_index: Dict[str, Set[str]] = {}
        self.file_pattern_index: Dict[str, Set[str]] = {}

        self._load_agents()

    def _load_agents(self):
        """Load all agents from the agents directory"""
        logger.info(f"Loading agents from {self.agents_dir}")

        # Standard contains-studio structure
        categories = [
            'engineering', 'design', 'product', 'marketing',
            'project-management', 'studio-operations', 'testing', 'bonus'
        ]

        for category in categories:
            category_path = os.path.join(self.agents_dir, category)
            if os.path.exists(category_path):
                self._load_category(category, category_path)

        # Also scan for individual .md files
        for root, dirs, files in os.walk(self.agents_dir):
            for file in files:
                if file.endswith('.md') and file != 'README.md':
                    self._load_agent_file(os.path.join(root, file))

        logger.info(f"Loaded {len(self.agents)} agents")

    def _load_category(self, category: str, category_path: str):
        """Load all agents from a category directory"""
        for file in os.listdir(category_path):
            if file.endswith('.md'):
                agent_path = os.path.join(category_path, file)
                self._load_agent_file(agent_path)

    def _load_agent_file(self, file_path: str):
        """Parse and load an agent definition from a markdown file"""
        try:
            with open(file_path, 'r') as f:
                content = f.read()

            # Parse YAML frontmatter
            frontmatter, body = self._parse_frontmatter(content)

            if not frontmatter.get('name'):
                return

            # Extract agent info
            name = frontmatter['name']
            category = self._infer_category(file_path)
            description = frontmatter.get('description', '')

            # Extract keywords from description and examples
            keywords = self._extract_keywords(description)

            # Determine trigger types
            trigger_types = self._determine_trigger_types(frontmatter, body)

            # Extract file patterns
            file_patterns = self._extract_file_patterns(body)

            # Get tools
            tools = frontmatter.get('tools', [])
            if isinstance(tools, str):
                tools = [t.strip() for t in tools.split(',')]

            # Extract examples
            examples = self._extract_examples(body)

            # Create capability
            capability = AgentCapability(
                name=name,
                category=category,
                description=description,
                keywords=keywords,
                trigger_types=trigger_types,
                file_patterns=file_patterns,
                tools=tools,
                examples=examples,
                priority=self._calculate_priority(category, tools)
            )

            self.agents[name] = capability

            # Build indexes
            for keyword in keywords:
                if keyword not in self.keyword_index:
                    self.keyword_index[keyword] = set()
                self.keyword_index[keyword].add(name)

            for pattern in file_patterns:
                if pattern not in self.file_pattern_index:
                    self.file_pattern_index[pattern] = set()
                self.file_pattern_index[pattern].add(name)

        except Exception as e:
            logger.warning(f"Error loading agent from {file_path}: {e}")

    def _parse_frontmatter(self, content: str) -> tuple:
        """Parse YAML frontmatter from markdown"""
        if not content.startswith('---'):
            return {}, content

        # Find end of frontmatter
        end = content.find('---', 4)
        if end == -1:
            return {}, content

        frontmatter_str = content[4:end].strip()
        body = content[end + 4:].strip()

        # Simple YAML parsing
        frontmatter = {}
        for line in frontmatter_str.split('\n'):
            if ':' in line:
                key, value = line.split(':', 1)
                frontmatter[key.strip()] = value.strip()

        return frontmatter, body

    def _infer_category(self, file_path: str) -> AgentCategory:
        """Infer category from file path"""
        path_lower = file_path.lower()

        if 'engineering' in path_lower:
            return AgentCategory.ENGINEERING
        elif 'design' in path_lower:
            return AgentCategory.DESIGN
        elif 'product' in path_lower:
            return AgentCategory.PRODUCT
        elif 'marketing' in path_lower:
            return AgentCategory.MARKETING
        elif 'project-management' in path_lower or 'project' in path_lower:
            return AgentCategory.PROJECT_MANAGEMENT
        elif 'studio-operations' in path_lower or 'operations' in path_lower:
            return AgentCategory.STUDIO_OPERATIONS
        elif 'testing' in path_lower or 'test' in path_lower:
            return AgentCategory.TESTING
        else:
            return AgentCategory.BONUS

    def _extract_keywords(self, description: str) -> List[str]:
        """Extract keywords from description"""
        keywords = []

        # Common tech keywords
        tech_keywords = [
            'ai', 'ml', 'api', 'backend', 'frontend', 'mobile', 'ios', 'android',
            'react', 'vue', 'svelte', 'angular', 'typescript', 'javascript',
            'python', 'rust', 'go', 'java', 'swift', 'kotlin',
            'ui', 'ux', 'design', 'component', 'layout', 'style',
            'test', 'testing', 'unit', 'integration', 'e2e',
            'deploy', 'ci', 'cd', 'docker', 'kubernetes',
            'database', 'sql', 'nosql', 'redis', 'postgres',
            'auth', 'authentication', 'oauth', 'jwt',
            'payment', 'stripe', 'billing',
            'design', 'figma', 'mockup', 'prototype',
            'marketing', 'seo', 'social', 'content',
            'analytics', 'metrics', 'data',
            'performance', 'optimization', 'speed',
            'security', 'compliance', 'legal',
            'documentation', 'docs', 'readme'
        ]

        description_lower = description.lower()

        # Extract mentioned tech keywords
        for keyword in tech_keywords:
            if keyword in description_lower:
                keywords.append(keyword)

        # Extract from examples
        example_keywords = re.findall(r'example>\nContext: ([^\n]+)', description_lower)
        keywords.extend(example_keywords)

        # Extract action verbs
        actions = ['build', 'create', 'design', 'implement', 'refactor', 'test',
                   'deploy', 'optimize', 'fix', 'add', 'integrate', 'setup']
        for action in actions:
            if action in description_lower:
                keywords.append(action)

        return list(set(keywords))

    def _determine_trigger_types(self, frontmatter: Dict, body: str) -> List[TriggerType]:
        """Determine how this agent can be triggered"""
        trigger_types = [TriggerType.EXPLICIT, TriggerType.KEYWORD]

        body_lower = body.lower()

        # Check for proactive triggers
        if 'proactively' in body_lower or 'trigger automatically' in body_lower:
            trigger_types.append(TriggerType.PROACTIVE)

        # Check for file pattern triggers
        if any(ext in body_lower for ext in ['.tsx', '.py', '.rs', '.go', '.java']):
            trigger_types.append(TriggerType.FILE_PATTERN)

        # Check for context triggers
        if 'context' in body_lower or 'when' in body_lower:
            trigger_types.append(TriggerType.CONTEXT)

        return trigger_types

    def _extract_file_patterns(self, body: str) -> List[str]:
        """Extract file patterns that trigger this agent"""
        patterns = []

        # Common patterns
        extensions = re.findall(r'\.([a-z]+)', body)
        for ext in set(extensions):
            if len(ext) <= 5:  # Reasonable file extension
                patterns.append(f'*.{ext}')

        # Path patterns
        paths = re.findall(r'([a-z-]+/)', body.lower())
        for path in set(paths):
            if path in ['src/', 'components/', 'tests/', 'docs/', 'api/']:
                patterns.append(path)

        return patterns

    def _extract_examples(self, body: str) -> List[Dict]:
        """Extract usage examples"""
        examples = []

        # Find example blocks
        example_blocks = re.findall(r'<example>(.*?)</example>', body, re.DOTALL)

        for block in example_blocks:
            context_match = re.search(r'Context: ([^\n]+)', block)
            user_match = re.search(r'user: "([^"]+)"', block)
            assistant_match = re.search(r'assistant: "([^"]+)"', block)

            if context_match and user_match:
                examples.append({
                    'context': context_match.group(1),
                    'user_request': user_match.group(1),
                    'response': assistant_match.group(1) if assistant_match else '',
                    'full_block': block.strip()
                })

        return examples

    def _calculate_priority(self, category: AgentCategory, tools: List[str]) -> int:
        """Calculate agent priority for selection"""
        priority = 5

        # Engineering agents tend to be higher priority
        if category == AgentCategory.ENGINEERING:
            priority = 7
        elif category == AgentCategory.DESIGN:
            priority = 6
        elif category == AgentCategory.TESTING:
            priority = 8  # Testing is proactive

        # Boost agents with more tools
        priority += min(len(tools), 3)

        return min(priority, 10)

    def find_agents_by_keywords(self, text: str) -> List[tuple]:
        """Find agents matching keywords in text, sorted by relevance"""
        text_lower = text.lower()
        words = set(text_lower.split())

        matches = []

        for agent_name, agent in self.agents.items():
            score = 0

            # Check keyword matches
            for keyword in agent.keywords:
                if keyword.lower() in text_lower:
                    score += 1

            # Check examples
            for example in agent.examples:
                if example['user_request'].lower() in text_lower:
                    score += 3

            # Check direct name mention
            if agent_name.lower() in text_lower:
                score += 10

            if score > 0:
                matches.append((agent_name, score, agent))

        # Sort by score, then priority
        matches.sort(key=lambda x: (x[1], x[2].priority), reverse=True)

        return matches

    def find_agents_by_files(self, files: List[str]) -> List[tuple]:
        """Find agents that should handle specific file types"""
        matches = []

        for file_path in files:
            file_lower = file_path.lower()

            for agent_name, agent in self.agents.items():
                score = 0

                # Check file patterns
                for pattern in agent.file_patterns:
                    if pattern in file_lower:
                        score += 1

                if score > 0:
                    matches.append((agent_name, score, agent))

        matches.sort(key=lambda x: (x[1], x[2].priority), reverse=True)
        return matches

    def find_proactive_agents(self, context: Dict) -> List[str]:
        """Find agents that should trigger proactively"""
        proactive = []

        for agent_name, agent in self.agents.items():
            if TriggerType.PROACTIVE in agent.trigger_types:
                # Check context
                if self._check_proactive_context(agent, context):
                    proactive.append(agent_name)

        return proactive

    def _check_proactive_context(self, agent: AgentCapability, context: Dict) -> bool:
        """Check if agent should trigger proactively in this context"""
        # Test-writer-fixer triggers after code changes
        if agent.name == 'test-writer-fixer':
            return context.get('code_modified', False)

        # Whimsy-injector triggers after UI changes
        if agent.name == 'whimsy-injector':
            return context.get('ui_modified', False)

        # Studio-coach triggers on complex tasks
        if agent.name == 'studio-coach':
            return context.get('complexity', 0) > 7

        return False

    def get_agent(self, name: str) -> Optional[AgentCapability]:
        """Get agent by name"""
        return self.agents.get(name)

    def get_all_agents(self) -> Dict[str, AgentCapability]:
        """Get all registered agents"""
        return self.agents

    def get_agents_by_category(self, category: AgentCategory) -> List[AgentCapability]:
        """Get all agents in a category"""
        return [a for a in self.agents.values() if a.category == category]


# Pre-configured agent mappings for contains-studio agents
CONTAINS_STUDIO_AGENTS = {
    # Engineering
    'ai-engineer': {
        'keywords': ['ai', 'ml', 'llm', 'machine learning', 'recommendation', 'chatbot', 'computer vision'],
        'triggers': ['implement ai', 'add ml', 'integrate llm', 'build recommendation']
    },
    'backend-architect': {
        'keywords': ['api', 'backend', 'server', 'database', 'microservices'],
        'triggers': ['design api', 'build backend', 'create database schema']
    },
    'devops-automator': {
        'keywords': ['deploy', 'ci/cd', 'docker', 'kubernetes', 'infrastructure'],
        'triggers': ['set up deployment', 'configure ci', 'deploy to production']
    },
    'frontend-developer': {
        'keywords': ['frontend', 'ui', 'component', 'react', 'vue', 'svelte'],
        'triggers': ['build component', 'create ui', 'implement frontend']
    },
    'mobile-app-builder': {
        'keywords': ['mobile', 'ios', 'android', 'react native', 'swift', 'kotlin'],
        'triggers': ['build mobile app', 'create ios', 'develop android']
    },
    'rapid-prototyper': {
        'keywords': ['mvp', 'prototype', 'quick', 'scaffold', 'new app'],
        'triggers': ['create prototype', 'build mvp', 'scaffold project', 'new app idea']
    },
    'test-writer-fixer': {
        'keywords': ['test', 'testing', 'coverage'],
        'triggers': ['write tests', 'add coverage', 'test this'],
        'proactive': True
    },

    # Design
    'brand-guardian': {
        'keywords': ['brand', 'logo', 'identity', 'guidelines'],
        'triggers': ['design brand', 'create logo', 'brand guidelines']
    },
    'ui-designer': {
        'keywords': ['ui', 'interface', 'design', 'component design'],
        'triggers': ['design ui', 'create interface', 'ui design']
    },
    'ux-researcher': {
        'keywords': ['ux', 'user research', 'usability', 'user experience'],
        'triggers': ['user research', 'ux study', 'usability test']
    },
    'visual-storyteller': {
        'keywords': ['visual', 'story', 'graphic', 'illustration'],
        'triggers': ['create visual', 'design graphics', 'story telling']
    },
    'whimsy-injector': {
        'keywords': ['delight', 'surprise', 'fun', 'animation'],
        'triggers': ['add delight', 'make fun', 'surprise users'],
        'proactive': True
    },

    # Product
    'feedback-synthesizer': {
        'keywords': ['feedback', 'reviews', 'complaints', 'user input'],
        'triggers': ['analyze feedback', 'synthesize reviews', 'user complaints']
    },
    'sprint-prioritizer': {
        'keywords': ['sprint', 'priority', 'roadmap', 'planning'],
        'triggers': ['plan sprint', 'prioritize features', 'sprint planning']
    },
    'trend-researcher': {
        'keywords': ['trend', 'viral', 'market research', 'opportunity'],
        'triggers': ['research trends', 'whats trending', 'market analysis']
    },

    # Marketing
    'app-store-optimizer': {
        'keywords': ['app store', 'aso', 'store listing', 'keywords'],
        'triggers': ['optimize app store', 'improve aso', 'store listing']
    },
    'content-creator': {
        'keywords': ['content', 'blog', 'social media', 'copy'],
        'triggers': ['create content', 'write blog', 'social content']
    },
    'growth-hacker': {
        'keywords': ['growth', 'viral', 'acquisition', 'funnel'],
        'triggers': ['growth strategy', 'viral loop', 'acquisition']
    },
    'tiktok-strategist': {
        'keywords': ['tiktok', 'video', 'viral', 'content'],
        'triggers': ['tiktok strategy', 'viral video', 'tiktok content']
    },

    # Project Management
    'experiment-tracker': {
        'keywords': ['experiment', 'a/b test', 'feature flag'],
        'triggers': ['track experiment', 'a/b testing', 'feature flags'],
        'proactive': True
    },
    'project-shipper': {
        'keywords': ['launch', 'ship', 'release', 'deploy'],
        'triggers': ['prepare launch', 'ship project', 'release management']
    },
    'studio-producer': {
        'keywords': ['coordinate', 'team', 'workflow', 'manage'],
        'triggers': ['coordinate team', 'manage project', 'workflow']
    },

    # Testing
    'api-tester': {
        'keywords': ['api test', 'load test', 'endpoint testing'],
        'triggers': ['test api', 'load testing', 'endpoint test']
    },
    'performance-benchmarker': {
        'keywords': ['performance', 'benchmark', 'speed', 'optimization'],
        'triggers': ['benchmark performance', 'speed test', 'optimize']
    },
    'test-results-analyzer': {
        'keywords': ['test results', 'analyze tests', 'test failures'],
        'triggers': ['analyze test results', 'test failures', 'test report']
    },

    # Studio Operations
    'analytics-reporter': {
        'keywords': ['analytics', 'metrics', 'data', 'reports'],
        'triggers': ['generate report', 'analyze metrics', 'analytics']
    },
    'finance-tracker': {
        'keywords': ['finance', 'budget', 'costs', 'revenue'],
        'triggers': ['track costs', 'budget analysis', 'financial report']
    },
    'infrastructure-maintainer': {
        'keywords': ['infrastructure', 'servers', 'monitoring', 'uptime'],
        'triggers': ['check infrastructure', 'server health', 'monitoring']
    },
    'support-responder': {
        'keywords': ['support', 'help', 'customer service'],
        'triggers': ['handle support', 'customer inquiry', 'help ticket']
    },

    # Bonus
    'joker': {
        'keywords': ['joke', 'humor', 'funny', 'laugh'],
        'triggers': ['tell joke', 'add humor', 'make funny']
    },
    'studio-coach': {
        'keywords': ['coach', 'guidance', 'help', 'advice'],
        'triggers': ['need help', 'guidance', 'coach'],
        'proactive': True
    }
}