"""
Config Aggregator Collector - Org-wide resource discovery via AWS Config.

COS1-02: HITL runs `runbooks inventory collect --config-aggregator` and discovers
resources across all org accounts in <10s (vs 60s×67 accounts per-account loop).

Uses the organization-aggregator Config Aggregator to query AWS::EC2::Instance,
AWS::EC2::Volume, AWS::S3::Bucket, and other resource types via SQL-like queries.
"""

import json
import logging
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any, Dict, List, Optional

import boto3
from botocore.exceptions import ClientError

from runbooks.common.rich_utils import console

logger = logging.getLogger(__name__)

_CONFIG_SEMAPHORE = threading.Semaphore(4)

CONFIG_RESOURCE_TYPES = {
    "ec2": "AWS::EC2::Instance",
    "ebs": "AWS::EC2::Volume",
    "s3": "AWS::S3::Bucket",
    "rds": "AWS::RDS::DBInstance",
    "lambda": "AWS::Lambda::Function",
    "vpc": "AWS::EC2::VPC",
    "iam_role": "AWS::IAM::Role",
    "cloudformation": "AWS::CloudFormation::Stack",
    "acm": "AWS::ACM::Certificate",
}


def detect_aggregator(session: boto3.Session, region: str = "ap-southeast-2") -> Optional[str]:
    """Detect the first available Config Aggregator in the account."""
    try:
        from runbooks.inventory.cross_validator import CROSS_VALIDATE_CLIENT_CONFIG

        config_client = session.client("config", region_name=region, config=CROSS_VALIDATE_CLIENT_CONFIG)
        response = config_client.describe_configuration_aggregators()
        aggregators = response.get("ConfigurationAggregators", [])
        if aggregators:
            name = aggregators[0]["ConfigurationAggregatorName"]
            logger.info(f"Found Config Aggregator: {name}")
            return name
        return None
    except ClientError as e:
        logger.debug(f"Config Aggregator detection failed: {e}")
        return None


def query_config_resources(
    session: boto3.Session,
    aggregator_name: str,
    resource_type: str,
    region: str = "ap-southeast-2",
) -> List[Dict[str, Any]]:
    """Query Config Aggregator for resources of a given type using SQL expression."""
    from runbooks.inventory.cross_validator import CROSS_VALIDATE_CLIENT_CONFIG

    config_client = session.client("config", region_name=region, config=CROSS_VALIDATE_CLIENT_CONFIG)
    query = (
        f"SELECT resourceId, resourceName, resourceType, accountId, awsRegion, configuration "
        f"WHERE resourceType = '{resource_type}'"
    )

    resources = []
    next_token = None

    while True:
        params = {
            "ConfigurationAggregatorName": aggregator_name,
            "Expression": query,
            "MaxResults": 100,
        }
        if next_token:
            params["NextToken"] = next_token

        try:
            response = config_client.select_aggregate_resource_config(**params)
        except ClientError as e:
            logger.warning(f"Config query failed for {resource_type}: {e}")
            break

        for result_str in response.get("Results", []):
            try:
                resources.append(json.loads(result_str))
            except json.JSONDecodeError:
                pass

        next_token = response.get("NextToken")
        if not next_token:
            break

    return resources


def normalize_config_result(resource: Dict[str, Any], resource_category: str) -> Dict[str, Any]:
    """Normalize a Config Aggregator result into standard inventory format."""
    return {
        "resource_id": resource.get("resourceId", ""),
        "resource_name": resource.get("resourceName", ""),
        "resource_type": resource_category,
        "account_id": resource.get("accountId", ""),
        "region": resource.get("awsRegion", ""),
        "configuration": resource.get("configuration", {}),
    }


def _query_category(
    session: boto3.Session,
    aggregator_name: str,
    category: str,
    region: str,
) -> Dict[str, Any]:
    """Worker for parallel Config Aggregator queries. Never raises."""
    aws_type = CONFIG_RESOURCE_TYPES.get(category)
    if not aws_type:
        return {"category": category, "resources": [], "error": f"Unknown: {category}"}
    with _CONFIG_SEMAPHORE:
        try:
            raw = query_config_resources(session, aggregator_name, aws_type, region)
            normalized = [normalize_config_result(r, category) for r in raw]
            return {"category": category, "resources": normalized, "error": None}
        except Exception as e:
            logger.warning(f"Failed to query {category}: {e}")
            return {"category": category, "resources": [], "error": str(e)}


def collect_via_aggregator(
    session: boto3.Session,
    aggregator_name: str,
    resource_types: Optional[List[str]] = None,
    region: str = "ap-southeast-2",
    max_workers: int = 4,
) -> Dict[str, Any]:
    """
    Collect inventory via Config Aggregator for org-wide discovery.

    Returns dict with 'resources' (list) and 'summary' (per-type counts).
    """
    start_time = time.time()
    types_to_query = list(CONFIG_RESOURCE_TYPES.keys()) if resource_types is None else resource_types
    all_resources: List[Dict[str, Any]] = []
    summary: Dict[str, int] = {}

    valid_types = [c for c in types_to_query if CONFIG_RESOURCE_TYPES.get(c)]
    for inv in [c for c in types_to_query if not CONFIG_RESOURCE_TYPES.get(c)]:
        logger.warning(f"Unknown resource category: {inv}")

    console.print(
        f"[dim]Starting parallel Config Aggregator collection "
        f"({len(valid_types)} types, {max_workers} workers)...[/dim]"
    )

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {executor.submit(_query_category, session, aggregator_name, cat, region): cat for cat in valid_types}
        for future in as_completed(futures):
            result = future.result()
            all_resources.extend(result["resources"])
            summary[result["category"].upper()] = len(result["resources"])
            if result["error"]:
                logger.warning(f"Config Aggregator error: {result['error']}")

    duration = round(time.time() - start_time, 2)
    total_accounts = len(set(r["account_id"] for r in all_resources if r["account_id"]))
    console.print(f"[dim]Config Aggregator collection completed in {duration}s[/dim]")

    return {
        "resources": all_resources,
        "summary": summary,
        "total_resources": len(all_resources),
        "total_accounts": total_accounts,
        "duration_seconds": duration,
        "aggregator_name": aggregator_name,
        "resource_types_queried": types_to_query,
    }
