"""
Multi-Account Remediation Engine - Enterprise AWS Organizations Support

## Overview

This module provides enterprise-grade multi-account remediation capabilities,
enabling bulk security and compliance operations across AWS Organizations.
Enhanced and migrated from the original references/bulk_run.py with enterprise
features including SSO integration, parallel execution, and comprehensive auditing.

## Original Functionality Enhanced

Migrated and enhanced from references/bulk_run.py with these improvements:
- Enterprise SSO integration with AWS IAM Identity Center
- Parallel execution across multiple accounts
- Comprehensive error handling and retry logic
- Progress tracking and status reporting
- Integration with remediation base classes
- Compliance evidence aggregation across accounts

## Key Features

- **SSO Integration**: Seamless AWS SSO credential management
- **Parallel Execution**: Concurrent operations across multiple accounts
- **Progress Tracking**: Real-time status and progress reporting
- **Error Handling**: Comprehensive retry logic and failure management
- **Audit Trail**: Complete operation tracking across all accounts
- **Compliance Aggregation**: Consolidated compliance evidence

## Example Usage

```python
from runbooks.remediation import MultiAccountRemediator
from runbooks.inventory.models.account import AWSAccount

# Initialize multi-account remediator
remediator = MultiAccountRemediator(
    sso_start_url="https://d-123456789.awsapps.com/start",
    parallel_execution=True,
    max_workers=5
)

# Define target accounts using dynamic discovery
# Example: Get accounts from AWS Organizations or environment configuration
accounts = get_accounts_from_environment() or discover_organization_accounts()

# Execute bulk S3 security remediation
results = remediator.bulk_s3_security(
    accounts=accounts,
    operations=["block_public_access", "enforce_ssl"],
    dry_run=False
)
```

Version: 0.7.8 - Enterprise Production Ready
"""

import concurrent.futures
import json
import os
import time
import uuid
import webbrowser
from typing import Any, Dict, List, Optional, Set, Tuple

import botocore.exceptions
from botocore.exceptions import ClientError
from loguru import logger

from runbooks.common.aws_client_factory import create_sso_client, create_sso_oidc_client

from runbooks.inventory.models.account import AWSAccount
from runbooks.remediation.base import BaseRemediation, RemediationContext, RemediationResult, RemediationStatus
from runbooks.common.profile_utils import create_management_session
from runbooks.remediation.universal_account_discovery import (
    UniversalAccountDiscovery,
    AWSAccount as UniversalAWSAccount,
)


class MultiAccountRemediator:
    """
    Enterprise multi-account remediation orchestrator.

    Provides comprehensive capabilities for executing remediation operations
    across multiple AWS accounts with enterprise features including SSO
    integration, parallel execution, and consolidated reporting.

    ## Key Features

    - **SSO Integration**: AWS IAM Identity Center support
    - **Parallel Execution**: Concurrent operations with configurable workers
    - **Progress Tracking**: Real-time status and completion tracking
    - **Error Handling**: Comprehensive retry and failure management
    - **Compliance Aggregation**: Consolidated compliance evidence
    - **Audit Trail**: Complete operation tracking across accounts

    ## Example Usage

    ```python
    # Initialize with SSO configuration
    remediator = MultiAccountRemediator(
        sso_start_url="https://d-123456789.awsapps.com/start",
        role_name="CloudOpsRemediationRole",
        parallel_execution=True
    )

    # Execute bulk remediation
    results = remediator.execute_bulk_operation(
        accounts=accounts,
        operation_class="S3SecurityRemediation",
        operation_method="block_public_access",
        operation_kwargs={"bucket_name": "critical-bucket"}
    )
    ```
    """

    def __init__(
        self,
        sso_start_url: Optional[str] = None,
        role_name: str = "power-user",
        parallel_execution: bool = True,
        max_workers: int = 5,
        **kwargs,
    ):
        """
        Initialize multi-account remediator.

        Args:
            sso_start_url: AWS SSO start URL for authentication
            role_name: IAM role name for cross-account access
            parallel_execution: Enable parallel execution across accounts
            max_workers: Maximum number of concurrent workers
            **kwargs: Additional configuration parameters
        """
        self.sso_start_url = sso_start_url or os.getenv("ACCESS_PORTAL_URL")
        self.role_name = role_name
        self.parallel_execution = parallel_execution
        self.max_workers = max_workers

        # Enterprise configuration
        self.retry_attempts = kwargs.get("retry_attempts", 3)
        self.retry_delay = kwargs.get("retry_delay", 5)
        self.timeout_seconds = kwargs.get("timeout_seconds", 300)

        # Progress tracking
        self.operation_id = str(uuid.uuid4())
        self.progress_callback = kwargs.get("progress_callback")

        # Credential cache
        self._credentials_cache = {}
        self._sso_token = None

        logger.info(f"MultiAccountRemediator initialized with operation ID: {self.operation_id}")

    def get_all_sso_credentials(self) -> Dict[str, Dict[str, str]]:
        """
        Get AWS SSO credentials for all accessible accounts.

        Enhanced from original bulk_run.py with improved error handling,
        credential caching, and progress tracking.

        Returns:
            Dictionary mapping account profiles to credentials
        """
        if not self.sso_start_url:
            raise ValueError("SSO start URL is required for multi-account operations")

        credentials = {}

        try:
            # Create SSO OIDC client
            sso_oidc = create_sso_oidc_client(region=os.getenv("AWS_DEFAULT_REGION", "ap-southeast-2"))

            # Register client
            client_creds = sso_oidc.register_client(clientName="CloudOpsRemediation", clientType="public")

            # Get device authorization
            device_auth = sso_oidc.start_device_authorization(
                clientId=client_creds["clientId"],
                clientSecret=client_creds["clientSecret"],
                startUrl=self.sso_start_url,
            )

            logger.info(f"Please authorize access at: {device_auth['verificationUriComplete']}")
            logger.info(f"User code: {device_auth['userCode']}")

            # Open browser automatically
            try:
                webbrowser.open(device_auth["verificationUriComplete"])
            except Exception as e:
                logger.warning(f"Could not open browser automatically: {e}")

            # Wait for user authorization
            token = None
            max_retries = 60  # 5 minutes with 5-second intervals
            retry_count = 0

            while not token and retry_count < max_retries:
                try:
                    token = sso_oidc.create_token(
                        clientId=client_creds["clientId"],
                        clientSecret=client_creds["clientSecret"],
                        grantType="urn:ietf:params:oauth:grant-type:device_code",
                        deviceCode=device_auth["deviceCode"],
                    )
                    self._sso_token = token

                except sso_oidc.exceptions.AuthorizationPendingException:
                    logger.info("Waiting for authorization... Please complete the process in your browser.")
                    time.sleep(5)
                    retry_count += 1
                except Exception as e:
                    logger.error(f"Authorization error: {e}")
                    break

            if not token:
                raise Exception("SSO authorization timed out or failed")

            # Create SSO client
            sso = create_sso_client(region="ap-southeast-2")

            # List all accounts with pagination
            all_accounts = []
            paginator = sso.get_paginator("list_accounts")
            for page in paginator.paginate(accessToken=token["accessToken"]):
                all_accounts.extend(page["accountList"])

            logger.info(f"Found {len(all_accounts)} accessible accounts")

            # Get credentials for each account
            for account in all_accounts:
                account_id = account["accountId"]

                # Get available roles for the account
                roles = []
                role_paginator = sso.get_paginator("list_account_roles")
                for role_page in role_paginator.paginate(accessToken=token["accessToken"], accountId=account_id):
                    roles.extend(role_page["roleList"])

                # Find the specified role
                target_role = None
                for role in roles:
                    if role["roleName"] == self.role_name:
                        target_role = role
                        break

                if target_role:
                    try:
                        # Get temporary credentials
                        creds = sso.get_role_credentials(
                            accessToken=token["accessToken"], accountId=account_id, roleName=target_role["roleName"]
                        )

                        profile_key = f"{account_id}_{target_role['roleName']}"
                        credentials[profile_key] = {
                            "aws_access_key_id": creds["roleCredentials"]["accessKeyId"],
                            "aws_secret_access_key": creds["roleCredentials"]["secretAccessKey"],
                            "aws_session_token": creds["roleCredentials"]["sessionToken"],
                            "account_id": account_id,
                            "account_name": account.get("accountName", account_id),
                            "role_name": target_role["roleName"],
                        }

                        logger.debug(f"Retrieved credentials for account {account_id}")

                    except Exception as e:
                        logger.warning(f"Failed to get credentials for account {account_id}: {e}")
                else:
                    logger.warning(f"Role {self.role_name} not found in account {account_id}")

            logger.info(f"Successfully retrieved credentials for {len(credentials)} accounts")
            self._credentials_cache = credentials

            return credentials

        except Exception as e:
            logger.error(f"Failed to retrieve SSO credentials: {e}")
            raise

    def get_credentials_for_account(self, account_id: str) -> Optional[Dict[str, str]]:
        """
        Get credentials for a specific account.

        Args:
            account_id: AWS account ID

        Returns:
            Credentials dictionary or None if not found
        """
        # Check cache first
        for profile_key, creds in self._credentials_cache.items():
            if creds.get("account_id") == account_id:
                return creds

        # If not in cache, refresh credentials
        if not self._credentials_cache:
            self.get_all_sso_credentials()

        # Check again after refresh
        for profile_key, creds in self._credentials_cache.items():
            if creds.get("account_id") == account_id:
                return creds

        logger.warning(f"No credentials found for account {account_id}")
        return None

    def execute_single_account_operation(
        self,
        account: AWSAccount,
        operation_class: str,
        operation_method: str,
        operation_kwargs: Dict[str, Any],
        **context_kwargs,
    ) -> List[RemediationResult]:
        """
        Execute remediation operation on a single account.

        Args:
            account: Target AWS account
            operation_class: Remediation class name (e.g., "S3SecurityRemediation")
            operation_method: Method name to execute
            operation_kwargs: Arguments for the operation method
            **context_kwargs: Additional context parameters

        Returns:
            List of remediation results for the account
        """
        account_results = []

        try:
            # Get credentials for the account
            credentials = self.get_credentials_for_account(account.account_id)
            if not credentials:
                raise Exception(f"No credentials available for account {account.account_id}")

            # Set environment variables for AWS SDK
            original_env = {}
            for key in ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN"]:
                original_env[key] = os.environ.get(key)

            os.environ["AWS_ACCESS_KEY_ID"] = credentials["aws_access_key_id"]
            os.environ["AWS_SECRET_ACCESS_KEY"] = credentials["aws_secret_access_key"]
            os.environ["AWS_SESSION_TOKEN"] = credentials["aws_session_token"]

            try:
                # Create remediation context
                context = RemediationContext(
                    account=account, operation_type=f"{operation_class}.{operation_method}", **context_kwargs
                )

                # Dynamically import and instantiate the remediation class
                if operation_class == "S3SecurityRemediation":
                    from runbooks.remediation.s3_remediation import S3SecurityRemediation

                    remediation_instance = S3SecurityRemediation()
                else:
                    raise ValueError(f"Unsupported remediation class: {operation_class}")

                # Execute the operation
                method = getattr(remediation_instance, operation_method)
                operation_results = method(context, **operation_kwargs)

                account_results.extend(operation_results)

                logger.info(f"Completed operation {operation_method} for account {account.account_id}")

            finally:
                # Restore original environment variables
                for key, value in original_env.items():
                    if value is None:
                        os.environ.pop(key, None)
                    else:
                        os.environ[key] = value

        except Exception as e:
            logger.error(f"Failed to execute operation {operation_method} for account {account.account_id}: {e}")

            # Create error result
            error_context = RemediationContext(
                account=account, operation_type=f"{operation_class}.{operation_method}", **context_kwargs
            )

            error_result = RemediationResult(
                context=error_context, status=RemediationStatus.FAILED, error_message=str(e)
            )
            account_results.append(error_result)

        return account_results

    def execute_bulk_operation(
        self,
        accounts: List[AWSAccount],
        operation_class: str,
        operation_method: str,
        operation_kwargs: Dict[str, Any],
        **context_kwargs,
    ) -> List[RemediationResult]:
        """
        Execute remediation operation across multiple accounts.

        Args:
            accounts: List of target AWS accounts
            operation_class: Remediation class name
            operation_method: Method name to execute
            operation_kwargs: Arguments for the operation method
            **context_kwargs: Additional context parameters

        Returns:
            Consolidated list of remediation results
        """
        logger.info(f"Starting bulk operation {operation_method} across {len(accounts)} accounts")

        all_results = []

        if self.parallel_execution and len(accounts) > 1:
            # Parallel execution
            with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                # Submit all account operations
                future_to_account = {
                    executor.submit(
                        self.execute_single_account_operation,
                        account,
                        operation_class,
                        operation_method,
                        operation_kwargs,
                        **context_kwargs,
                    ): account
                    for account in accounts
                }

                # Collect results as they complete
                for future in concurrent.futures.as_completed(future_to_account, timeout=self.timeout_seconds):
                    account = future_to_account[future]
                    try:
                        account_results = future.result()
                        all_results.extend(account_results)

                        # Progress callback
                        if self.progress_callback:
                            self.progress_callback(account, account_results)

                    except Exception as e:
                        logger.error(f"Parallel execution error for account {account.account_id}: {e}")
        else:
            # Sequential execution
            for account in accounts:
                account_results = self.execute_single_account_operation(
                    account, operation_class, operation_method, operation_kwargs, **context_kwargs
                )
                all_results.extend(account_results)

                # Progress callback
                if self.progress_callback:
                    self.progress_callback(account, account_results)

        # Generate summary statistics
        successful_results = [r for r in all_results if r.success]
        failed_results = [r for r in all_results if r.failed]

        logger.info(
            f"Bulk operation completed: {len(successful_results)} successful, "
            f"{len(failed_results)} failed across {len(accounts)} accounts"
        )

        return all_results

    def bulk_s3_security(
        self, accounts: List[AWSAccount], operations: List[str], bucket_patterns: Optional[List[str]] = None, **kwargs
    ) -> List[RemediationResult]:
        """
        Execute bulk S3 security operations across multiple accounts.

        Args:
            accounts: List of target AWS accounts
            operations: List of S3 security operations to execute
            bucket_patterns: Optional bucket name patterns to target
            **kwargs: Additional parameters

        Returns:
            Consolidated list of remediation results
        """
        logger.info(f"Starting bulk S3 security operations: {operations} across {len(accounts)} accounts")

        all_results = []

        for operation in operations:
            if operation == "block_public_access":
                operation_results = self.execute_bulk_operation(
                    accounts=accounts,
                    operation_class="S3SecurityRemediation",
                    operation_method="block_public_access",
                    operation_kwargs=kwargs,
                    **kwargs,
                )
            elif operation == "enforce_ssl":
                operation_results = self.execute_bulk_operation(
                    accounts=accounts,
                    operation_class="S3SecurityRemediation",
                    operation_method="enforce_ssl",
                    operation_kwargs=kwargs,
                    **kwargs,
                )
            elif operation == "enable_encryption":
                operation_results = self.execute_bulk_operation(
                    accounts=accounts,
                    operation_class="S3SecurityRemediation",
                    operation_method="enable_encryption",
                    operation_kwargs=kwargs,
                    **kwargs,
                )
            else:
                logger.warning(f"Unsupported S3 security operation: {operation}")
                continue

            all_results.extend(operation_results)

        return all_results

    def generate_compliance_report(self, results: List[RemediationResult]) -> Dict[str, Any]:
        """
        Generate consolidated compliance report from multi-account results.

        Args:
            results: List of remediation results from multiple accounts

        Returns:
            Comprehensive compliance report
        """
        # Group results by account
        results_by_account = {}
        for result in results:
            account_id = result.context.account.account_id
            if account_id not in results_by_account:
                results_by_account[account_id] = []
            results_by_account[account_id].append(result)

        # Generate compliance summary
        compliance_summary = {
            "total_accounts": len(results_by_account),
            "total_operations": len(results),
            "successful_operations": len([r for r in results if r.success]),
            "failed_operations": len([r for r in results if r.failed]),
            "compliance_frameworks": set(),
            "compliance_controls": set(),
            "accounts_summary": {},
        }

        # Account-level summaries
        for account_id, account_results in results_by_account.items():
            account_summary = {
                "account_id": account_id,
                "total_operations": len(account_results),
                "successful_operations": len([r for r in account_results if r.success]),
                "failed_operations": len([r for r in account_results if r.failed]),
                "compliance_controls": [],
            }

            # Aggregate compliance information
            for result in account_results:
                if hasattr(result.context, "compliance_mapping"):
                    mapping = result.context.compliance_mapping
                    compliance_summary["compliance_controls"].update(mapping.cis_controls)
                    compliance_summary["compliance_controls"].update(mapping.nist_categories)
                    account_summary["compliance_controls"].extend(mapping.cis_controls)

            compliance_summary["accounts_summary"][account_id] = account_summary

        # Convert sets to lists for JSON serialization
        compliance_summary["compliance_controls"] = list(compliance_summary["compliance_controls"])

        return compliance_summary


# Dynamic account discovery functions for enterprise security operations
def get_accounts_from_environment(profile: Optional[str] = None) -> Optional[List[AWSAccount]]:
    """
    Get AWS accounts using universal account discovery system.

    Uses enhanced discovery with support for:
    - Environment variables (REMEDIATION_TARGET_ACCOUNTS)
    - Configuration files (REMEDIATION_ACCOUNT_CONFIG)
    - AWS Organizations API (automatic discovery)
    - Current account fallback (single account mode)

    Args:
        profile: AWS profile to use for discovery

    Returns:
        List of AWSAccount objects or None if not configured
    """
    try:
        # Use universal account discovery system
        discovery = UniversalAccountDiscovery(profile=profile)
        universal_accounts = discovery.discover_target_accounts()

        if not universal_accounts:
            return None

        # Convert to legacy AWSAccount format for compatibility
        legacy_accounts = []
        for universal_account in universal_accounts:
            legacy_account = AWSAccount(
                universal_account.account_id,
                universal_account.account_name or f"account-{universal_account.account_id}",
            )
            legacy_accounts.append(legacy_account)

        logger.info(f"Using {len(legacy_accounts)} accounts discovered via universal discovery system")
        return legacy_accounts

    except Exception as e:
        logger.warning(f"Failed to discover accounts using universal discovery: {e}")
        return None


def discover_organization_accounts(profile: Optional[str] = None) -> List[AWSAccount]:
    """
    Discover AWS accounts using universal discovery system.

    Enhanced to use the universal account discovery system which provides:
    - Organizations API discovery (if available)
    - Environment variable fallback
    - Configuration file support
    - Current account fallback

    Args:
        profile: AWS profile for discovery (universal profile management)

    Returns:
        List of discovered AWSAccount objects
    """
    try:
        # Use universal account discovery system for Organizations discovery
        discovery = UniversalAccountDiscovery(profile=profile)
        universal_accounts = discovery._get_accounts_from_organizations()

        if not universal_accounts:
            # Fallback to other discovery methods
            logger.info("Organizations API not available, trying other discovery methods...")
            universal_accounts = discovery.discover_target_accounts()

        # Convert to legacy AWSAccount format for compatibility
        legacy_accounts = []
        for universal_account in universal_accounts:
            if universal_account.status == "ACTIVE":
                legacy_account = AWSAccount(
                    universal_account.account_id,
                    universal_account.account_name or f"org-account-{universal_account.account_id}",
                )
                legacy_accounts.append(legacy_account)

        logger.info(f"Discovered {len(legacy_accounts)} active AWS accounts via universal discovery")
        return legacy_accounts

    except Exception as e:
        logger.warning(f"Failed to discover organization accounts: {e}")
        # Universal discovery handles all fallback scenarios
        return []


def _determine_account_environment(account_name: str) -> str:
    """
    Determine account environment based on account name patterns.

    Args:
        account_name: AWS account name

    Returns:
        Environment classification
    """
    name_lower = account_name.lower()

    # Common environment patterns
    if any(env in name_lower for env in ["prod", "production"]):
        return "production"
    elif any(env in name_lower for env in ["staging", "stage", "uat"]):
        return "staging"
    elif any(env in name_lower for env in ["dev", "development"]):
        return "development"
    elif any(env in name_lower for env in ["test", "testing"]):
        return "testing"
    elif any(env in name_lower for env in ["sandbox", "sb"]):
        return "sandbox"
    else:
        return "unknown"
