kaizen-agentic/src/kaizen_agentic/extensions.py

"""Extension mechanisms for project-specific Kaizen agent customizations."""

import json
import yaml
from pathlib import Path
from typing import Dict, List, Optional, Any
from dataclasses import dataclass, field
from enum import Enum


class ExtensionType(Enum):
    """Types of agent extensions."""

    CONFIGURATION_OVERLAY = "config_overlay"  # Override default configurations
    FUNCTIONAL_EXTENSION = "functional_extension"  # Add new functionality
    WORKFLOW_INTEGRATION = "workflow_integration"  # Integrate with project workflows
    CUSTOM_COMMANDS = "custom_commands"  # Add custom commands
    DATA_TRANSFORMATION = "data_transformation"  # Transform data formats
    ENVIRONMENT_ADAPTATION = "env_adaptation"  # Adapt to specific environments


@dataclass
class AgentExtension:
    """Defines an extension to a Kaizen agent."""

    name: str
    base_agent: str  # The Kaizen agent this extends
    extension_type: ExtensionType
    description: str
    version: str = "1.0.0"
    author: Optional[str] = None

    # Extension configuration
    configuration: Dict[str, Any] = field(default_factory=dict)
    custom_commands: Dict[str, str] = field(default_factory=dict)
    workflow_hooks: Dict[str, str] = field(default_factory=dict)
    data_transformations: Dict[str, str] = field(default_factory=dict)
    environment_overrides: Dict[str, Any] = field(default_factory=dict)

    # Metadata
    dependencies: List[str] = field(default_factory=list)
    compatibility: List[str] = field(default_factory=list)  # Compatible Kaizen versions
    enabled: bool = True


@dataclass
class ProjectExtensionRegistry:
    """Registry of extensions for a project."""

    project_path: Path
    extensions: List[AgentExtension] = field(default_factory=list)
    global_config: Dict[str, Any] = field(default_factory=dict)


class ExtensionManager:
    """Manages agent extensions for projects."""

    def __init__(self, project_path: Path):
        self.project_path = Path(project_path)
        self.extensions_dir = self.project_path / ".kaizen" / "extensions"
        self.config_file = self.project_path / ".kaizen" / "extensions.yml"

    def create_extension(
        self,
        name: str,
        base_agent: str,
        extension_type: ExtensionType,
        description: str,
        **kwargs,
    ) -> AgentExtension:
        """Create a new agent extension."""
        extension = AgentExtension(
            name=name,
            base_agent=base_agent,
            extension_type=extension_type,
            description=description,
            **kwargs,
        )

        self._save_extension(extension)
        return extension

    def install_extension(self, extension_path: Path) -> AgentExtension:
        """Install an extension from a file."""
        if extension_path.suffix == ".json":
            with open(extension_path) as f:
                data = json.load(f)
        elif extension_path.suffix in [".yml", ".yaml"]:
            with open(extension_path) as f:
                data = yaml.safe_load(f)
        else:
            raise ValueError(
                f"Unsupported extension file format: {extension_path.suffix}"
            )

        extension = AgentExtension(**data)
        self._save_extension(extension)
        return extension

    def list_extensions(self, base_agent: Optional[str] = None) -> List[AgentExtension]:
        """List all extensions, optionally filtered by base agent."""
        extensions = self._load_extensions()

        if base_agent:
            return [ext for ext in extensions if ext.base_agent == base_agent]

        return extensions

    def get_extension(self, name: str) -> Optional[AgentExtension]:
        """Get a specific extension by name."""
        extensions = self._load_extensions()
        return next((ext for ext in extensions if ext.name == name), None)

    def enable_extension(self, name: str) -> bool:
        """Enable an extension."""
        return self._toggle_extension(name, True)

    def disable_extension(self, name: str) -> bool:
        """Disable an extension."""
        return self._toggle_extension(name, False)

    def remove_extension(self, name: str) -> bool:
        """Remove an extension."""
        extensions = self._load_extensions()
        original_count = len(extensions)
        extensions = [ext for ext in extensions if ext.name != name]

        if len(extensions) < original_count:
            self._save_extensions(extensions)

            # Remove extension files
            extension_dir = self.extensions_dir / name
            if extension_dir.exists():
                import shutil

                shutil.rmtree(extension_dir)

            return True

        return False

    def get_effective_config(self, base_agent: str) -> Dict[str, Any]:
        """Get the effective configuration for an agent with all extensions applied."""
        base_config = self._get_base_agent_config(base_agent)
        extensions = [
            ext
            for ext in self._load_extensions()
            if ext.base_agent == base_agent and ext.enabled
        ]

        # Apply extensions in order
        for extension in extensions:
            base_config = self._apply_extension_config(base_config, extension)

        return base_config

    def create_project_specific_agent(
        self,
        name: str,
        base_agent: str,
        custom_instructions: str,
        custom_commands: Optional[Dict[str, str]] = None,
        environment_config: Optional[Dict[str, Any]] = None,
    ) -> AgentExtension:
        """Create a project-specific agent based on a Kaizen agent."""

        # Create extension configuration
        config = {
            "custom_instructions": custom_instructions,
            "project_context": {
                "name": self.project_path.name,
                "path": str(self.project_path),
                "type": self._detect_project_type(),
            },
        }

        if environment_config:
            config["environment"] = environment_config

        extension = self.create_extension(
            name=name,
            base_agent=base_agent,
            extension_type=ExtensionType.FUNCTIONAL_EXTENSION,
            description=f"Project-specific extension of {base_agent} for {self.project_path.name}",
            configuration=config,
            custom_commands=custom_commands or {},
            environment_overrides=environment_config or {},
        )

        # Create agent file
        self._create_extended_agent_file(extension)

        return extension

    def integrate_legacy_agent(
        self,
        legacy_agent_path: Path,
        target_kaizen_agent: str,
        migration_strategy: str = "preserve_functionality",
    ) -> AgentExtension:
        """Integrate a legacy agent as an extension to a Kaizen agent."""

        # Analyze legacy agent
        legacy_analysis = self._analyze_legacy_agent(legacy_agent_path)

        # Create extension based on analysis
        extension_name = f"{legacy_agent_path.stem}_integration"

        config = {
            "legacy_source": str(legacy_agent_path),
            "migration_strategy": migration_strategy,
            "preserved_functionality": legacy_analysis.get("functionality", []),
            "custom_config": legacy_analysis.get("config", {}),
        }

        extension = self.create_extension(
            name=extension_name,
            base_agent=target_kaizen_agent,
            extension_type=ExtensionType.WORKFLOW_INTEGRATION,
            description=f"Legacy integration of {legacy_agent_path.name}",
            configuration=config,
        )

        # Create migration wrapper
        self._create_migration_wrapper(extension, legacy_agent_path)

        return extension

    def _save_extension(self, extension: AgentExtension):
        """Save an extension to the registry."""
        extensions = self._load_extensions()

        # Remove existing extension with same name
        extensions = [ext for ext in extensions if ext.name != extension.name]
        extensions.append(extension)

        self._save_extensions(extensions)

        # Create extension directory and files
        extension_dir = self.extensions_dir / extension.name
        extension_dir.mkdir(parents=True, exist_ok=True)

        # Save extension definition
        with open(extension_dir / "extension.yml", "w") as f:
            # Convert dataclass to dict for YAML serialization
            data = {
                "name": extension.name,
                "base_agent": extension.base_agent,
                "extension_type": extension.extension_type.value,
                "description": extension.description,
                "version": extension.version,
                "author": extension.author,
                "configuration": extension.configuration,
                "custom_commands": extension.custom_commands,
                "workflow_hooks": extension.workflow_hooks,
                "data_transformations": extension.data_transformations,
                "environment_overrides": extension.environment_overrides,
                "dependencies": extension.dependencies,
                "compatibility": extension.compatibility,
                "enabled": extension.enabled,
            }
            yaml.dump(data, f, default_flow_style=False)

    def _load_extensions(self) -> List[AgentExtension]:
        """Load all extensions from the registry."""
        if not self.config_file.exists():
            return []

        try:
            with open(self.config_file) as f:
                data = yaml.safe_load(f) or {}

            extensions = []
            for ext_data in data.get("extensions", []):
                # Convert string back to enum
                ext_data["extension_type"] = ExtensionType(ext_data["extension_type"])
                extensions.append(AgentExtension(**ext_data))

            return extensions
        except Exception:
            return []

    def _save_extensions(self, extensions: List[AgentExtension]):
        """Save extensions to the registry."""
        self.config_file.parent.mkdir(parents=True, exist_ok=True)

        # Convert to serializable format
        data = {
            "extensions": [
                {
                    "name": ext.name,
                    "base_agent": ext.base_agent,
                    "extension_type": ext.extension_type.value,
                    "description": ext.description,
                    "version": ext.version,
                    "author": ext.author,
                    "configuration": ext.configuration,
                    "custom_commands": ext.custom_commands,
                    "workflow_hooks": ext.workflow_hooks,
                    "data_transformations": ext.data_transformations,
                    "environment_overrides": ext.environment_overrides,
                    "dependencies": ext.dependencies,
                    "compatibility": ext.compatibility,
                    "enabled": ext.enabled,
                }
                for ext in extensions
            ]
        }

        with open(self.config_file, "w") as f:
            yaml.dump(data, f, default_flow_style=False)

    def _toggle_extension(self, name: str, enabled: bool) -> bool:
        """Enable or disable an extension."""
        extensions = self._load_extensions()

        for extension in extensions:
            if extension.name == name:
                extension.enabled = enabled
                self._save_extensions(extensions)
                return True

        return False

    def _get_base_agent_config(self, base_agent: str) -> Dict[str, Any]:
        """Get the base configuration for a Kaizen agent."""
        # This would load the actual agent configuration
        # For now, return a basic structure
        return {
            "name": base_agent,
            "type": "kaizen_agent",
            "enabled": True,
            "config": {},
        }

    def _apply_extension_config(
        self, base_config: Dict[str, Any], extension: AgentExtension
    ) -> Dict[str, Any]:
        """Apply an extension's configuration to the base config."""
        config = base_config.copy()

        # Apply configuration overlays
        if extension.configuration:
            config["config"].update(extension.configuration)

        # Add custom commands
        if extension.custom_commands:
            config.setdefault("custom_commands", {}).update(extension.custom_commands)

        # Apply environment overrides
        if extension.environment_overrides:
            config.setdefault("environment", {}).update(extension.environment_overrides)

        # Add extension metadata
        config.setdefault("extensions", []).append(
            {
                "name": extension.name,
                "type": extension.extension_type.value,
                "version": extension.version,
            }
        )

        return config

    def _detect_project_type(self) -> str:
        """Detect the type of project."""
        if (self.project_path / "pyproject.toml").exists():
            return "python"
        elif (self.project_path / "package.json").exists():
            return "nodejs"
        elif (self.project_path / "Cargo.toml").exists():
            return "rust"
        elif (self.project_path / "go.mod").exists():
            return "go"
        else:
            return "generic"

    def _analyze_legacy_agent(self, agent_path: Path) -> Dict[str, Any]:
        """Analyze a legacy agent to understand its functionality."""
        analysis = {
            "functionality": [],
            "config": {},
            "commands": [],
            "dependencies": [],
        }

        if agent_path.suffix == ".py":
            # Analyze Python agent
            content = agent_path.read_text()

            # Simple analysis - look for class and function definitions
            import re

            classes = re.findall(r"class\s+(\w+)", content)
            functions = re.findall(r"def\s+(\w+)", content)

            analysis["functionality"] = classes + functions

        elif agent_path.suffix in [".yml", ".yaml"]:
            # Analyze YAML configuration
            try:
                with open(agent_path) as f:
                    data = yaml.safe_load(f)

                analysis["config"] = data
                analysis["commands"] = data.get("commands", [])

            except Exception:
                pass

        elif agent_path.suffix == ".json":
            # Analyze JSON configuration
            try:
                with open(agent_path) as f:
                    data = json.load(f)

                analysis["config"] = data
                analysis["functionality"] = data.get("features", [])

            except Exception:
                pass

        return analysis

    def _create_extended_agent_file(self, extension: AgentExtension):
        """Create an agent file for a project-specific extension."""
        agent_content = f"""---
name: {extension.name}
description: {extension.description}
base_agent: {extension.base_agent}
extension_type: {extension.extension_type.value}
version: {extension.version}
"""

        if extension.author:
            agent_content += f"author: {extension.author}\n"

        agent_content += "---\n\n"
        agent_content += f"# {extension.name}\n\n"
        agent_content += f"{extension.description}\n\n"
        agent_content += (
            f"This agent extends **{extension.base_agent}** "
            f"with project-specific functionality.\n\n"
        )

        if extension.configuration.get("custom_instructions"):
            agent_content += "## Custom Instructions\n\n"
            agent_content += f"{extension.configuration['custom_instructions']}\n\n"

        if extension.custom_commands:
            agent_content += "## Custom Commands\n\n"
            for cmd, desc in extension.custom_commands.items():
                agent_content += f"- **{cmd}**: {desc}\n"
            agent_content += "\n"

        # Save to agents directory
        agents_dir = self.project_path / "agents"
        agents_dir.mkdir(exist_ok=True)

        agent_file = agents_dir / f"agent-{extension.name}.md"
        agent_file.write_text(agent_content)

    def _create_migration_wrapper(self, extension: AgentExtension, legacy_path: Path):
        """Create a wrapper to integrate legacy agent functionality."""
        wrapper_dir = self.extensions_dir / extension.name
        wrapper_file = wrapper_dir / "legacy_wrapper.py"

        wrapper_content = f'''"""
Legacy agent wrapper for {extension.name}

This wrapper provides compatibility with the legacy agent at {legacy_path}
while integrating with the Kaizen agent {extension.base_agent}.
"""

import sys
from pathlib import Path

# Add legacy agent path to system path
legacy_path = Path("{legacy_path}").parent
if str(legacy_path) not in sys.path:
    sys.path.insert(0, str(legacy_path))

class LegacyWrapper:
    """Wrapper for legacy agent functionality."""

    def __init__(self):
        self.extension_config = {extension.configuration}
        self.legacy_path = Path("{legacy_path}")

    def execute_legacy_functionality(self, *args, **kwargs):
        """Execute legacy agent functionality."""
        # Implementation would depend on the legacy agent type
        pass

    def migrate_data(self, data):
        """Migrate data from legacy format to Kaizen format."""
        # Implementation for data transformation
        return data
'''

        wrapper_file.write_text(wrapper_content)


def create_extension_template(
    name: str, base_agent: str, project_path: Path, template_type: str = "basic"
) -> str:
    """Create a template for a new agent extension."""

    templates = {
        "basic": f"""# Extension Template: {name}

This template helps you create a custom extension for the {base_agent} agent.

## Configuration

```yaml
name: {name}
base_agent: {base_agent}
extension_type: functional_extension
description: "Custom extension for {base_agent}"
version: "1.0.0"
author: "Your Name"

configuration:
  custom_setting: "value"
  project_specific_config: {{}}

custom_commands:
  custom-command: "Description of custom command"

environment_overrides:
  CUSTOM_ENV_VAR: "value"
```

## Implementation

1. Define your custom configuration in the `configuration` section
2. Add custom commands in the `custom_commands` section
3. Override environment variables in `environment_overrides`
4. Implement any custom logic in separate Python files

## Usage

Save this as `.kaizen/extensions/{name}/extension.yml` and run:

```bash
kaizen-agentic extensions install {name}
```
""",
        "advanced": f"""# Advanced Extension Template: {name}

This template provides advanced customization options for the {base_agent} agent.

## Full Configuration

```yaml
name: {name}
base_agent: {base_agent}
extension_type: workflow_integration
description: "Advanced extension with workflow integration"
version: "1.0.0"
author: "Your Name"

configuration:
  # Custom instructions for the agent
  custom_instructions: |
    You are working on the {project_path.name} project.
    Follow these project-specific guidelines:
    - Use our coding standards
    - Reference our documentation style
    - Integrate with our CI/CD pipeline

  # Project context
  project_context:
    name: "{project_path.name}"
    type: "python"  # or nodejs, rust, etc.
    framework: "custom"

  # Custom behaviors
  behaviors:
    auto_commit: false
    generate_tests: true
    update_docs: true

custom_commands:
  project-setup: "Initialize project structure following our standards"
  custom-deploy: "Deploy using our specific deployment pipeline"
  generate-config: "Generate project-specific configuration files"

workflow_hooks:
  pre_action: "scripts/pre_action.py"
  post_action: "scripts/post_action.py"

data_transformations:
  input_format: "custom"
  output_format: "kaizen_standard"

environment_overrides:
  PROJECT_ROOT: "{project_path}"
  CUSTOM_CONFIG_PATH: "{project_path}/.config"

dependencies:
  - "requests>=2.25.0"
  - "pyyaml>=5.0.0"

compatibility:
  - "kaizen-agentic>=0.2.0"
```

## Custom Scripts

Create these files in `.kaizen/extensions/{name}/scripts/`:

- `pre_action.py` - Run before agent actions
- `post_action.py` - Run after agent actions
- `data_transform.py` - Transform data formats
- `validation.py` - Validate project-specific requirements

## Installation

1. Save the configuration as `.kaizen/extensions/{name}/extension.yml`
2. Add any custom scripts to the scripts directory
3. Install with: `kaizen-agentic extensions install {name}`
""",
    }

    return templates.get(template_type, templates["basic"])