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

import json
import yaml
from pathlib import Path
from typing import Dict, List, Optional, Any, Union
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}** 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"])