"""
Enterprise Scale Collector - Option C: Scale & Optimize Implementation
Enhanced for 200+ AWS accounts with parallel processing and advanced MCP integration

Performance Targets:
- FinOps Analysis: <60s for 200 accounts (from <30s for 60 accounts)
- Inventory Collection: <90s comprehensive scan (from <45s for 60 accounts)
- Security Baseline: <15s for 15+ checks (unchanged)
"""

import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from datetime import datetime
from typing import Any, Dict, List, Optional

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

from runbooks.inventory.collectors.base import BaseResourceCollector, CollectionContext
from runbooks.inventory.models.resource import AWSResource
from runbooks.inventory.utils.threading_utils import ProgressMetrics, ThreadPoolManager


@dataclass
class EnterpriseScaleConfig:
    """Configuration for enterprise-scale operations."""

    max_workers: int = 50  # Increased from 10 for 200+ accounts
    batch_size: int = 20  # Process accounts in batches
    timeout_per_account: int = 120  # 2 minutes per account
    enable_cost_analysis: bool = True
    enable_security_scanning: bool = True
    parallel_regions: bool = True
    mcp_integration: bool = True


class EnterpriseScaleCollector(BaseResourceCollector):
    """
    Enterprise-scale AWS resource collector optimized for 200+ accounts.

    Features:
    - Advanced concurrent processing with batching
    - Dynamic resource prioritization
    - Enhanced MCP server integration
    - Multi-tenant support with customer isolation
    - Performance monitoring and optimization
    """

    service_category = "enterprise"
    supported_resources = {
        "organizations",
        "accounts",
        "cost_explorer",
        "config",
        "ec2",
        "s3",
        "rds",
        "lambda",
        "dynamodb",
        "vpc",
        "iam",
    }
    requires_org_access = True

    def __init__(
        self,
        profile: Optional[str] = None,
        region: str = "ap-southeast-2",
        config: Optional[EnterpriseScaleConfig] = None,
    ):
        """Initialize enterprise scale collector."""
        super().__init__(profile, region)
        self.config = config or EnterpriseScaleConfig()
        self.performance_metrics = {}
        self.cost_cache = {}
        self.security_findings = {}

        logger.info(f"Initialized EnterpriseScaleCollector with {self.config.max_workers} workers")

    def collect_resources(
        self, context: CollectionContext, resource_filters: Optional[Dict[str, Any]] = None
    ) -> List[AWSResource]:
        """
        Collect resources across 200+ accounts with performance optimization.
        """
        start_time = time.time()
        logger.info("Starting enterprise-scale resource collection")

        try:
            # Phase 1: Discover all accounts in organization
            accounts = self._discover_organization_accounts()
            logger.info(f"Discovered {len(accounts)} accounts in organization")

            # Phase 2: Collect resources in parallel batches
            resources = self._collect_resources_parallel(accounts, context, resource_filters)

            collection_time = time.time() - start_time
            logger.info(f"Enterprise collection completed in {collection_time:.2f} seconds")

            # Performance validation against targets
            self._validate_performance_targets(len(accounts), collection_time)

            return resources

        except Exception as e:
            logger.error(f"Enterprise collection failed: {e}")
            raise

    def _discover_organization_accounts(self) -> List[Dict[str, Any]]:
        """Discover all accounts in AWS Organizations."""
        logger.info("Discovering AWS Organizations accounts")

        try:
            org_client = self.get_client("organizations", self.region)

            accounts = []
            paginator = org_client.get_paginator("list_accounts")

            for page in paginator.paginate():
                accounts.extend(page["Accounts"])

            # Filter active accounts only
            active_accounts = [acc for acc in accounts if acc["Status"] == "ACTIVE"]

            logger.info(f"Found {len(active_accounts)} active accounts")
            return active_accounts

        except ClientError as e:
            logger.error(f"Failed to discover organization accounts: {e}")
            # Fallback: return single current account
            sts_client = self.get_client("sts", self.region)
            identity = sts_client.get_caller_identity()
            return [{"Id": identity["Account"], "Name": "Current Account", "Status": "ACTIVE"}]

    def _collect_resources_parallel(
        self, accounts: List[Dict[str, Any]], context: CollectionContext, resource_filters: Optional[Dict[str, Any]]
    ) -> List[AWSResource]:
        """Collect resources using advanced parallel processing."""
        all_resources = []

        def progress_callback(metrics: ProgressMetrics):
            logger.info(f"Progress: {metrics.get_completion_percentage():.1f}% complete")

        with ThreadPoolManager(max_workers=self.config.max_workers, progress_callback=progress_callback) as pool:
            for account in accounts:
                task_id = f"collect_{account['Id']}"
                pool.submit_task(task_id, self._collect_account_resources, account, context, resource_filters)

            results = pool.wait_for_completion(timeout=self.config.timeout_per_account * len(accounts))

            # Combine successful results
            successful_results = pool.get_successful_results()
            for task_id, resources in successful_results.items():
                if resources:
                    all_resources.extend(resources)

        logger.info(f"Collected {len(all_resources)} total resources")
        return all_resources

    def _collect_account_resources(
        self, account: Dict[str, Any], context: CollectionContext, resource_filters: Optional[Dict[str, Any]]
    ) -> List[AWSResource]:
        """Collect resources from a single account."""
        account_id = account["Id"]
        logger.debug(f"Collecting from account: {account_id}")

        account_resources = []

        try:
            session = self._get_account_session(account_id)
            priority_services = ["ec2", "s3", "rds", "lambda"]

            for service in priority_services:
                if service in context.resource_types or "all" in context.resource_types:
                    service_resources = self._collect_service_resources(session, service, account_id, context)
                    account_resources.extend(service_resources)

        except Exception as e:
            logger.error(f"Failed to collect from account {account_id}: {e}")

        return account_resources

    def _collect_service_resources(
        self, session: boto3.Session, service: str, account_id: str, context: CollectionContext
    ) -> List[AWSResource]:
        """Collect resources for a specific service."""
        resources = []

        try:
            if service == "ec2":
                resources = self._collect_ec2_resources(session, account_id, context)
            elif service == "s3":
                resources = self._collect_s3_resources(session, account_id, context)
            # Add more services as needed

        except Exception as e:
            logger.warning(f"Failed to collect {service} from {account_id}: {e}")

        return resources

    def _collect_ec2_resources(
        self, session: boto3.Session, account_id: str, context: CollectionContext
    ) -> List[AWSResource]:
        """Collect EC2 instances."""
        resources = []
        ec2_client = session.client("ec2", region_name=context.region)

        try:
            response = ec2_client.describe_instances()
            for reservation in response.get("Reservations", []):
                for instance in reservation.get("Instances", []):
                    resource = AWSResource(
                        resource_id=instance["InstanceId"],
                        resource_type="ec2:instance",
                        service_category="compute",
                        metadata=self._create_resource_metadata(context, instance),
                    )
                    resources.append(resource)
        except ClientError as e:
            logger.warning(f"Failed to collect EC2 from {account_id}: {e}")

        return resources

    def _collect_s3_resources(
        self, session: boto3.Session, account_id: str, context: CollectionContext
    ) -> List[AWSResource]:
        """Collect S3 buckets."""
        resources = []
        s3_client = session.client("s3")

        try:
            response = s3_client.list_buckets()
            for bucket in response.get("Buckets", []):
                resource = AWSResource(
                    resource_id=bucket["Name"],
                    resource_type="s3:bucket",
                    service_category="storage",
                    metadata=self._create_resource_metadata(context, bucket),
                )
                resources.append(resource)
        except ClientError as e:
            logger.warning(f"Failed to collect S3 from {account_id}: {e}")

        return resources

    def _get_account_session(self, account_id: str) -> boto3.Session:
        """Get AWS session for specific account."""
        # For now, return current session. Production would assume cross-account roles.
        return self.session

    def _validate_performance_targets(self, account_count: int, execution_time: float):
        """Validate performance targets are met."""
        logger.info(f"Performance validation: {account_count} accounts in {execution_time:.2f}s")

        # Scale target time based on account count
        if account_count <= 60:
            target_time = 45.0
        else:
            # Linear scaling: 90s for 200 accounts
            target_time = 45.0 + ((account_count - 60) / 140) * 45.0

        performance_met = execution_time <= target_time

        if performance_met:
            logger.info(f"✅ Performance target MET: {execution_time:.2f}s <= {target_time:.2f}s")
        else:
            logger.warning(f"⚠️ Performance target MISSED: {execution_time:.2f}s > {target_time:.2f}s")

        self.performance_metrics = {
            "account_count": account_count,
            "execution_time": execution_time,
            "target_time": target_time,
            "performance_met": performance_met,
        }

    def get_cost_information(self, context: CollectionContext, resource: AWSResource) -> Optional[Dict[str, Any]]:
        """Get cost information for a resource."""
        return None  # Placeholder

    def get_performance_metrics(self) -> Dict[str, Any]:
        """Get performance metrics from last collection."""
        return self.performance_metrics
