"""
Key Management Command

This module provides commands for key management including generation,
validation, rotation, backup, and restoration.

Author: MCP Proxy Adapter Team
Version: 1.0.0
"""

import logging
import os
import shutil
from pathlib import Path
from datetime import datetime

from .base import Command
from .result import CommandResult, SuccessResult, ErrorResult
from ..core.certificate_utils import CertificateUtils

from mcp_proxy_adapter.core.logging import get_global_logger
logger = logging.getLogger(__name__)


class KeyResult:
    """
    Result class for key operations.

    Contains key information and operation status.
    """

    def __init__(
        self,
        key_path: str,
        key_type: str,
        key_size: int,
        created_date: Optional[str] = None,
        expiry_date: Optional[str] = None,
        status: str = "valid",
        error: Optional[str] = None,
    ):
        """
        Initialize key result.

        Args:
            key_path: Path to key file
            key_type: Type of key (RSA, ECDSA, etc.)
            key_size: Key size in bits
            created_date: Key creation date
            expiry_date: Key expiry date
            status: Key status (valid, expired, error)
            error: Error message if any
        """
        self.key_path = key_path
        self.key_type = key_type
        self.key_size = key_size
        self.created_date = created_date
        self.expiry_date = expiry_date
        self.status = status
        self.error = error

    def to_dict(self) -> Dict[str, Any]:
        """
        Convert to dictionary format.

        Returns:
            Dictionary representation
        """
        return {
            "key_path": self.key_path,
            "key_type": self.key_type,
            "key_size": self.key_size,
            "created_date": self.created_date,
            "expiry_date": self.expiry_date,
            "status": self.status,
            "error": self.error,
        }

class KeyManagementCommand(Command):
    """
    Command for key management.

    Provides methods for generating, validating, rotating, backing up, and restoring keys.
    """

    # Command metadata
    name = "key_management"
    version = "1.0.0"
    descr = "Private key generation, validation, and management"
    category = "security"
    author = "MCP Proxy Adapter Team"
    email = "team@mcp-proxy-adapter.com"
    source_url = "https://github.com/mcp-proxy-adapter"
    result_class = KeyResult

    def __init__(self):
        """Initialize key management command."""
        super().__init__()
        self.certificate_utils = CertificateUtils()

    async def execute(self, **kwargs) -> CommandResult:
        """
        Execute key management command.

        Args:
            **kwargs: Command parameters including:
                - action: Action to perform (key_generate, key_validate, key_rotate, key_backup, key_restore)
                - key_type: Type of key to generate (RSA, ECDSA)
                - key_size: Key size in bits for generation
                - output_path: Output path for key generation
                - password: Password for key encryption
                - key_path: Key file path for validation, rotation, backup
                - old_key_path: Old key path for rotation
                - new_key_path: New key path for rotation
                - cert_path: Certificate path for rotation
                - backup_old: Whether to backup old key during rotation
                - backup_path: Backup path for key backup/restore
                - encrypt_backup: Whether to encrypt backup

        Returns:
            CommandResult with key operation status
        """
        action = kwargs.get("action", "key_validate")

        if action == "key_generate":
            key_type = kwargs.get("key_type")
            key_size = kwargs.get("key_size")
            output_path = kwargs.get("output_path")
            password = kwargs.get("password")
            return await self.key_generate(key_type, key_size, output_path, password)
        elif action == "key_validate":
            key_path = kwargs.get("key_path")
            password = kwargs.get("password")
            return await self.key_validate(key_path, password)
        elif action == "key_rotate":
            old_key_path = kwargs.get("old_key_path")
            new_key_path = kwargs.get("new_key_path")
            cert_path = kwargs.get("cert_path")
            backup_old = kwargs.get("backup_old", True)
            return await self.key_rotate(
                old_key_path, new_key_path, cert_path, backup_old
            )
        elif action == "key_backup":
            key_path = kwargs.get("key_path")
            backup_path = kwargs.get("backup_path")
            encrypt_backup = kwargs.get("encrypt_backup", True)
            password = kwargs.get("password")
            return await self.key_backup(
                key_path, backup_path, encrypt_backup, password
            )
        elif action == "key_restore":
            backup_path = kwargs.get("backup_path")
            key_path = kwargs.get("key_path")
            password = kwargs.get("password")
            return await self.key_restore(backup_path, key_path, password)
        else:
            return ErrorResult(
                message=f"Unknown action: {action}. Supported actions: key_generate, key_validate, key_rotate, key_backup, key_restore"
            )

    async def key_generate(
        self,
        key_type: str,
        key_size: int,
        output_path: str,
        password: Optional[str] = None,
    ) -> CommandResult:
        """
        Generate a new private key.

        Args:
            key_type: Type of key to generate (RSA, ECDSA)
            key_size: Key size in bits
            output_path: Path to save the generated key
            password: Optional password to encrypt the key

        Returns:
            CommandResult with key generation status
        """
        try:
            get_global_logger().info(f"Generating {key_type} key with size {key_size} bits")

            # Validate parameters
            if key_type not in ["RSA", "ECDSA"]:
                return ErrorResult(message="Key type must be RSA or ECDSA")

            if key_type == "ECDSA":
                if key_size not in [256, 384, 521]:
                    return ErrorResult(
                        message="ECDSA key size must be 256, 384, or 521 bits"
                    )
            elif key_size < 1024:
                return ErrorResult(message="Key size must be at least 1024 bits")

            # Create output directory if it doesn't exist
            output_dir = os.path.dirname(output_path)
            if output_dir:
                Path(output_dir).mkdir(parents=True, exist_ok=True)

            # Generate key
            result = self.certificate_utils.generate_private_key(
                key_type, key_size, output_path
            )

            if not result["success"]:
                return ErrorResult(message=f"Key generation failed: {result['error']}")

            key_result = KeyResult(
                key_path=output_path,
                key_type=key_type,
                key_size=key_size,
                created_date=datetime.now().isoformat(),
                status="valid",
            )

            get_global_logger().info(f"Key generated successfully: {output_path}")
            return SuccessResult(data={"key": key_result.to_dict(), "details": result})

        except Exception as e:
            get_global_logger().error(f"Key generation failed: {e}")
            return ErrorResult(message=f"Key generation failed: {str(e)}")

    async def key_validate(
        self, key_path: str, password: Optional[str] = None
    ) -> CommandResult:
        """
        Validate a private key.

        Args:
            key_path: Path to key file to validate
            password: Optional password if key is encrypted

        Returns:
            CommandResult with key validation status
        """
        try:
            get_global_logger().info(f"Validating key: {key_path}")

            # Validate parameters
            if not key_path or not os.path.exists(key_path):
                return ErrorResult(message=f"Key file not found: {key_path}")

            # Validate key
            result = self.certificate_utils.validate_private_key(key_path)

            if not result["success"]:
                return ErrorResult(message=f"Key validation failed: {result['error']}")

            key_result = KeyResult(
                key_path=key_path,
                key_type=result.get("key_type", "unknown"),
                key_size=result.get("key_size", 0),
                created_date=result.get("created_date"),
                status="valid",
            )

            get_global_logger().info(f"Key validation completed: {key_path}")
            return SuccessResult(data={"key": key_result.to_dict()})

        except Exception as e:
            get_global_logger().error(f"Key validation failed: {e}")
            return ErrorResult(message=f"Key validation failed: {str(e)}")

    async def key_rotate(
        self,
        old_key_path: str,
        new_key_path: str,
        cert_path: Optional[str] = None,
        backup_old: bool = True,
    ) -> CommandResult:
        """
        Rotate a private key.

        Args:
            old_key_path: Path to old key file
            new_key_path: Path to new key file
            cert_path: Optional certificate path to update with new key
            backup_old: Whether to backup the old key

        Returns:
            CommandResult with key rotation status
        """
        try:
            get_global_logger().info(f"Rotating key from {old_key_path} to {new_key_path}")

            # Validate parameters
            if not old_key_path or not os.path.exists(old_key_path):
                return ErrorResult(message=f"Old key file not found: {old_key_path}")

            if not new_key_path or not os.path.exists(new_key_path):
                return ErrorResult(message=f"New key file not found: {new_key_path}")

            # Validate both keys
            old_key_validation = await self.key_validate(old_key_path)
            new_key_validation = await self.key_validate(new_key_path)

            if not old_key_validation.to_dict()["success"]:
                return ErrorResult(
                    message=f"Old key validation failed: {old_key_validation.to_dict()['error']['message']}"
                )

            if not new_key_validation.to_dict()["success"]:
                return ErrorResult(
                    message=f"New key validation failed: {new_key_validation.to_dict()['error']['message']}"
                )

            # Backup old key if requested
            backup_path = None
            if backup_old:
                backup_path = (
                    f"{old_key_path}.backup.{datetime.now().strftime('%Y%m%d_%H%M%S')}"
                )
                shutil.copy2(old_key_path, backup_path)
                get_global_logger().info(f"Old key backed up to: {backup_path}")

            # Update certificate if provided
            cert_updated = False
            if cert_path and os.path.exists(cert_path):
                try:
                    # This would require implementing certificate re-signing
                    # For now, we'll just note that it needs to be done
                    cert_updated = True
                    get_global_logger().info(
                        f"Certificate {cert_path} needs to be re-signed with new key"
                    )
                except Exception as e:
                    get_global_logger().warning(f"Could not update certificate {cert_path}: {e}")

            # Replace old key with new key
            shutil.copy2(new_key_path, old_key_path)

            get_global_logger().info(f"Key rotation completed successfully")
            return SuccessResult(
                data={
                    "old_key_path": old_key_path,
                    "new_key_path": new_key_path,
                    "backup_path": backup_path,
                    "cert_updated": cert_updated,
                    "message": "Key rotation completed successfully",
                }
            )

        except Exception as e:
            get_global_logger().error(f"Key rotation failed: {e}")
            return ErrorResult(message=f"Key rotation failed: {str(e)}")

    async def key_backup(
        self,
        key_path: str,
        backup_path: str,
        encrypt_backup: bool = True,
        password: Optional[str] = None,
    ) -> CommandResult:
        """
        Backup a private key.

        Args:
            key_path: Path to key file to backup
            backup_path: Path to save the backup
            encrypt_backup: Whether to encrypt the backup
            password: Password for backup encryption

        Returns:
            CommandResult with backup status
        """
        try:
            get_global_logger().info(f"Backing up key: {key_path}")

            # Validate parameters
            if not key_path or not os.path.exists(key_path):
                return ErrorResult(message=f"Key file not found: {key_path}")

            # Validate key before backup
            key_validation = await self.key_validate(key_path)
            if not key_validation.to_dict()["success"]:
                return ErrorResult(
                    message=f"Key validation failed before backup: {key_validation.to_dict()['error']['message']}"
                )

            # Create backup directory if it doesn't exist
            backup_dir = os.path.dirname(backup_path)
            if backup_dir:
                Path(backup_dir).mkdir(parents=True, exist_ok=True)

            # Create backup
            if encrypt_backup and password:
                # Encrypted backup
                result = self.certificate_utils.create_encrypted_backup(
                    key_path, backup_path, password
                )
                if not result["success"]:
                    return ErrorResult(
                        message=f"Encrypted backup failed: {result['error']}"
                    )
            else:
                # Simple file copy
                shutil.copy2(key_path, backup_path)

            # Verify backup
            if not os.path.exists(backup_path):
                return ErrorResult(message="Backup file was not created")

            get_global_logger().info(f"Key backup completed successfully: {backup_path}")
            return SuccessResult(
                data={
                    "key_path": key_path,
                    "backup_path": backup_path,
                    "encrypted": encrypt_backup,
                    "backup_size": os.path.getsize(backup_path),
                    "backup_date": datetime.now().isoformat(),
                }
            )

        except Exception as e:
            get_global_logger().error(f"Key backup failed: {e}")
            return ErrorResult(message=f"Key backup failed: {str(e)}")

    async def key_restore(
        self, backup_path: str, key_path: str, password: Optional[str] = None
    ) -> CommandResult:
        """
        Restore a private key from backup.

        Args:
            backup_path: Path to backup file
            key_path: Path to restore the key to
            password: Password if backup is encrypted

        Returns:
            CommandResult with restore status
        """
        try:
            get_global_logger().info(f"Restoring key from backup: {backup_path}")

            # Validate parameters
            if not backup_path or not os.path.exists(backup_path):
                return ErrorResult(message=f"Backup file not found: {backup_path}")

            # Create target directory if it doesn't exist
            key_dir = os.path.dirname(key_path)
            if key_dir:
                Path(key_dir).mkdir(parents=True, exist_ok=True)

            # Restore key
            if password:
                # Try encrypted restore first
                result = self.certificate_utils.restore_encrypted_backup(
                    backup_path, key_path, password
                )
                if not result["success"]:
                    return ErrorResult(
                        message=f"Encrypted restore failed: {result['error']}"
                    )
            else:
                # Simple file copy
                shutil.copy2(backup_path, key_path)

            # Validate restored key
            key_validation = await self.key_validate(key_path)
            if not key_validation.to_dict()["success"]:
                return ErrorResult(
                    message=f"Restored key validation failed: {key_validation.to_dict()['error']['message']}"
                )

            get_global_logger().info(f"Key restore completed successfully: {key_path}")
            return SuccessResult(
                data={
                    "backup_path": backup_path,
                    "key_path": key_path,
                    "restore_date": datetime.now().isoformat(),
                    "key_info": (
                        key_validation.to_dict().get("data", {}).get("key")
                        if key_validation.to_dict().get("success")
                        else None
                    ),
                }
            )

        except Exception as e:
            get_global_logger().error(f"Key restore failed: {e}")
            return ErrorResult(message=f"Key restore failed: {str(e)}")
