#!/usr/bin/env python3
"""
Claude Code Subagent Management System

This script provides comprehensive management for Claude Code subagents,
including initialization, coordination, and workflow execution.
"""

import asyncio
import json
import logging
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import Any

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class SubagentStatus(Enum):
    AVAILABLE = "available"
    BUSY = "busy"
    ERROR = "error"
    OFFLINE = "offline"


@dataclass
class SubagentInstance:
    name: str
    config: dict[str, Any]
    status: SubagentStatus = SubagentStatus.AVAILABLE
    current_task: str | None = None
    load_factor: float = 0.0


class SubagentManager:
    """Advanced subagent management and coordination system"""

    def __init__(self, registry_path: str = ".claude-code/subagents/registry.json"):
        self.registry_path = Path(registry_path)
        self.subagents: dict[str, SubagentInstance] = {}
        self.collaboration_patterns: dict[str, list[str]] = {}
        self.active_workflows: dict[str, dict] = {}

        self._load_registry()
        self._initialize_subagents()

    def _load_registry(self):
        """Load subagent registry configuration"""
        try:
            with self.registry_path.open() as f:
                self.registry = json.load(f)

            self.collaboration_patterns = self.registry.get("interaction_patterns", {})
            logger.info(f"Loaded registry with {len(self.registry.get('subagents', {}))} subagents")

        except FileNotFoundError:
            logger.error(f"Registry file not found: {self.registry_path}")
            self.registry = {"subagents": {}}
        except json.JSONDecodeError as e:
            logger.error(f"Invalid JSON in registry: {e}")
            self.registry = {"subagents": {}}

    def _initialize_subagents(self):
        """Initialize all subagents from registry"""
        for _category, subagents in self.registry.get("subagents", {}).items():
            for name, config in subagents.items():
                try:
                    subagent_config = self._load_subagent_config(config["file"])
                    self.subagents[name] = SubagentInstance(name=name, config=subagent_config)
                    logger.info(f"Initialized subagent: {name}")
                except Exception as e:
                    logger.error(f"Failed to initialize subagent {name}: {e}")

    def _load_subagent_config(self, config_file: str) -> dict[str, Any]:
        """Load individual subagent configuration"""
        config_path = self.registry_path.parent / config_file

        try:
            with config_path.open() as f:
                return json.load(f)
        except FileNotFoundError:
            logger.warning(f"Config file not found: {config_path}")
            return {}
        except json.JSONDecodeError as e:
            logger.error(f"Invalid JSON in config {config_path}: {e}")
            return {}

    def get_subagent_capabilities(self, subagent_name: str) -> list[str]:
        """Get capabilities of a specific subagent"""
        if subagent_name in self.subagents:
            return self.subagents[subagent_name].config.get("capabilities", [])
        return []

    def find_subagents_by_capability(self, capability: str) -> list[str]:
        """Find all subagents with a specific capability"""
        matching_subagents = []
        for name, instance in self.subagents.items():
            if capability in instance.config.get("capabilities", []):
                matching_subagents.append(name)
        return matching_subagents

    def get_collaboration_pattern(self, pattern_name: str) -> list[str]:
        """Get predefined collaboration pattern"""
        return self.collaboration_patterns.get(pattern_name, [])

    def analyze_task_requirements(self, task_description: str) -> dict[str, Any]:
        """Analyze task and determine required capabilities and subagents"""
        # This would use NLP/ML in a real implementation
        # For now, provide a rule-based analysis

        capabilities_needed = []
        priority_subagents = []

        # Code quality tasks
        if any(keyword in task_description.lower() for keyword in ["lint", "format", "quality", "ruff", "type"]):
            capabilities_needed.extend(["static-code-analysis", "type-checking"])
            priority_subagents.append("python-quality-analyst")

        # Testing tasks
        if any(keyword in task_description.lower() for keyword in ["test", "bdd", "behave", "coverage"]):
            capabilities_needed.extend(["bdd-scenario-design", "test-execution"])
            priority_subagents.extend(["test-architect", "test-executor"])

        # Security tasks
        if any(keyword in task_description.lower() for keyword in ["security", "vulnerability", "scan", "audit"]):
            capabilities_needed.extend(["security-scanning", "vulnerability-assessment"])
            priority_subagents.append("security-auditor")

        # Performance tasks
        if any(keyword in task_description.lower() for keyword in ["performance", "optimize", "profile", "benchmark"]):
            capabilities_needed.extend(["performance-profiling", "optimization"])
            priority_subagents.append("performance-optimizer")

        # Container/deployment tasks
        if any(keyword in task_description.lower() for keyword in ["docker", "container", "deploy", "kubernetes"]):
            capabilities_needed.extend(["container-design", "deployment"])
            priority_subagents.extend(["container-architect", "kubernetes-specialist"])

        return {
            "capabilities_needed": capabilities_needed,
            "priority_subagents": priority_subagents,
            "complexity": self._assess_task_complexity(task_description),
            "estimated_duration": self._estimate_task_duration(task_description),
        }

    def _assess_task_complexity(self, task_description: str) -> str:
        """Assess task complexity based on description"""
        complexity_indicators = {
            "simple": ["fix", "update", "change"],
            "medium": ["implement", "create", "design", "optimize"],
            "complex": ["architect", "integrate", "migrate", "refactor"],
            "very_complex": ["orchestrate", "coordinate", "comprehensive"],
        }

        task_lower = task_description.lower()
        for complexity, indicators in complexity_indicators.items():
            if any(indicator in task_lower for indicator in indicators):
                return complexity

        return "medium"  # default

    def _estimate_task_duration(self, task_description: str) -> str:
        """Estimate task duration based on complexity and scope"""
        task_lower = task_description.lower()

        if any(keyword in task_lower for keyword in ["quick", "simple", "minor"]):
            return "short"  # < 30 minutes
        elif any(keyword in task_lower for keyword in ["comprehensive", "complete", "full"]):
            return "long"  # > 2 hours
        else:
            return "medium"  # 30 minutes - 2 hours

    async def execute_workflow(self, workflow_name: str, context: dict[str, Any]) -> dict[str, Any]:
        """Execute a predefined workflow pattern"""
        if workflow_name not in self.collaboration_patterns:
            raise ValueError(f"Unknown workflow: {workflow_name}")

        workflow_id = f"{workflow_name}_{asyncio.get_event_loop().time()}"
        self.active_workflows[workflow_id] = {
            "name": workflow_name,
            "status": "running",
            "context": context,
            "start_time": asyncio.get_event_loop().time(),
            "subagents": self.collaboration_patterns[workflow_name],
            "results": {},
        }

        try:
            results = await self._execute_workflow_steps(workflow_id)
            self.active_workflows[workflow_id]["status"] = "completed"
            self.active_workflows[workflow_id]["results"] = results
            return results

        except Exception as e:
            self.active_workflows[workflow_id]["status"] = "failed"
            self.active_workflows[workflow_id]["error"] = str(e)
            logger.error(f"Workflow {workflow_name} failed: {e}")
            raise

    async def _execute_workflow_steps(self, workflow_id: str) -> dict[str, Any]:
        """Execute workflow steps with proper coordination"""
        workflow = self.active_workflows[workflow_id]
        subagent_names = workflow["subagents"]
        context = workflow["context"]

        results = {}

        # For now, execute sequentially - in real implementation,
        # this would use intelligent parallel execution
        for subagent_name in subagent_names:
            if subagent_name in self.subagents:
                subagent = self.subagents[subagent_name]

                # Update subagent status
                subagent.status = SubagentStatus.BUSY
                subagent.current_task = workflow_id

                try:
                    # Simulate subagent execution
                    result = await self._simulate_subagent_execution(subagent_name, context)
                    results[subagent_name] = result

                    # Update context with results for next subagent
                    context.update({f"{subagent_name}_result": result})

                finally:
                    # Reset subagent status
                    subagent.status = SubagentStatus.AVAILABLE
                    subagent.current_task = None

        return results

    async def _simulate_subagent_execution(self, subagent_name: str, _context: dict[str, Any]) -> dict[str, Any]:
        """Simulate subagent execution (placeholder for actual implementation)"""
        subagent = self.subagents[subagent_name]

        # Simulate processing time based on subagent complexity
        processing_time = len(subagent.config.get("capabilities", [])) * 0.5
        await asyncio.sleep(processing_time)

        return {
            "status": "completed",
            "subagent": subagent_name,
            "capabilities_used": subagent.config.get("capabilities", []),
            "processing_time": processing_time,
            "recommendations": f"Recommendations from {subagent_name}",
            "artifacts": [f"artifact_{subagent_name}.json"],
        }

    def get_system_status(self) -> dict[str, Any]:
        """Get comprehensive system status"""
        status_summary = {
            "total_subagents": len(self.subagents),
            "available_subagents": sum(1 for s in self.subagents.values() if s.status == SubagentStatus.AVAILABLE),
            "busy_subagents": sum(1 for s in self.subagents.values() if s.status == SubagentStatus.BUSY),
            "active_workflows": len([w for w in self.active_workflows.values() if w["status"] == "running"]),
            "subagent_details": {},
        }

        for name, instance in self.subagents.items():
            status_summary["subagent_details"][name] = {
                "status": instance.status.value,
                "current_task": instance.current_task,
                "capabilities": len(instance.config.get("capabilities", [])),
                "load_factor": instance.load_factor,
            }

        return status_summary

    def get_available_workflows(self) -> list[str]:
        """Get list of available workflow patterns"""
        return list(self.collaboration_patterns.keys())

    def recommend_subagents(self, task_description: str) -> dict[str, Any]:
        """Recommend optimal subagents for a given task"""
        analysis = self.analyze_task_requirements(task_description)

        # Find subagents with required capabilities
        recommended_subagents = []
        for capability in analysis["capabilities_needed"]:
            capable_subagents = self.find_subagents_by_capability(capability)
            recommended_subagents.extend(capable_subagents)

        # Remove duplicates and prioritize
        recommended_subagents = list(set(recommended_subagents))

        # Add priority subagents
        for priority_subagent in analysis["priority_subagents"]:
            if priority_subagent not in recommended_subagents:
                recommended_subagents.append(priority_subagent)

        # Check availability
        available_subagents = [
            name
            for name in recommended_subagents
            if name in self.subagents and self.subagents[name].status == SubagentStatus.AVAILABLE
        ]

        return {
            "task_analysis": analysis,
            "recommended_subagents": recommended_subagents,
            "available_subagents": available_subagents,
            "coordination_recommendation": self._recommend_coordination_pattern(recommended_subagents),
        }

    def _recommend_coordination_pattern(self, subagents: list[str]) -> dict[str, Any]:
        """Recommend coordination pattern based on subagent combination"""
        # Find matching workflow patterns
        matching_patterns = []
        for pattern_name, pattern_subagents in self.collaboration_patterns.items():
            overlap = set(subagents) & set(pattern_subagents)
            if overlap:
                matching_patterns.append(
                    {
                        "pattern": pattern_name,
                        "overlap_count": len(overlap),
                        "coverage": len(overlap) / len(set(subagents) | set(pattern_subagents)),
                    }
                )

        # Sort by coverage
        matching_patterns.sort(key=lambda x: x["coverage"], reverse=True)

        return {
            "recommended_pattern": matching_patterns[0]["pattern"] if matching_patterns else None,
            "pattern_options": matching_patterns,
            "custom_coordination_needed": len(matching_patterns) == 0,
        }


# CLI Interface
def main():
    """Command-line interface for subagent management"""
    import argparse

    parser = argparse.ArgumentParser(description="Claude Code Subagent Manager")
    parser.add_argument("--status", action="store_true", help="Show system status")
    parser.add_argument("--workflows", action="store_true", help="List available workflows")
    parser.add_argument("--recommend", type=str, help="Recommend subagents for task")
    parser.add_argument("--execute", type=str, help="Execute workflow")

    args = parser.parse_args()

    manager = SubagentManager()

    if args.status:
        status = manager.get_system_status()
        print(json.dumps(status, indent=2))

    elif args.workflows:
        workflows = manager.get_available_workflows()
        print("Available workflows:")
        for workflow in workflows:
            print(f"  - {workflow}")

    elif args.recommend:
        recommendations = manager.recommend_subagents(args.recommend)
        print(json.dumps(recommendations, indent=2))

    elif args.execute:

        async def run_workflow():
            result = await manager.execute_workflow(args.execute, {})
            print(json.dumps(result, indent=2))

        asyncio.run(run_workflow())

    else:
        parser.print_help()


if __name__ == "__main__":
    main()