from __future__ import annotations

import pprint
import sys
from collections import defaultdict
from pathlib import Path
from typing import Callable, DefaultDict

from cloudenv.server.providers.aws.operations_data import AWS_SERVICE_OPERATIONS

# -----------------------------------------------------------------------------
# Paths / import setup (match other scripts in this repo)
# -----------------------------------------------------------------------------

ROOT = Path(__file__).resolve().parents[1]  # .../cloudenv
AWS_DIR = ROOT / "server" / "providers" / "aws"

# Allows running via `python cloudenv/scripts/...` from repo root.
sys.path.insert(0, str(ROOT.parent))


CORE_SERVICES: set[str] = {
    "s3",
    "ec2",
    "sns",
    "sqs",
    "lambda",
    "apigateway",
    "apigatewayv2",
    "cloudwatch",
    "logs",
    "kinesis",
    "dynamodb",
}


def _generic_bucket(op: str) -> str:
    read_prefixes = ("List", "Get", "Describe", "Head")
    write_prefixes = (
        "Put",
        "Create",
        "Update",
        "Modify",
        "Associate",
        "Attach",
        "Enable",
        "Tag",
        "Add",
        "Start",
        "Send",
        "Publish",
        "Upload",
        "Register",
        "Grant",
        "Link",
        "Set",
        "Import",
    )
    delete_prefixes = (
        "Delete",
        "Remove",
        "Terminate",
        "Detach",
        "Disassociate",
        "Untag",
        "Disable",
        "Stop",
        "Revoke",
        "Deregister",
        "Unsubscribe",
        "Abort",
        "Cancel",
    )

    if op.startswith(read_prefixes):
        return "read"
    if op.startswith(delete_prefixes):
        return "delete"
    if op.startswith(write_prefixes):
        return "write"
    return "misc"


# -----------------------------------------------------------------------------
# Core service classifiers (must return a suffix for every op)
# -----------------------------------------------------------------------------


def classify_s3(op: str) -> str:
    # Bucket lifecycle
    if op in ("CreateBucket", "DeleteBucket"):
        return "manage_bucket"

    # Bucket list / discovery
    if op in ("ListBuckets",):
        return "list_buckets"
    if op in ("ListObjects", "ListObjectsV2", "ListObjectVersions"):
        return "list_bucket"

    # Head / Get bucket
    if op in ("HeadBucket",):
        return "read_bucket"

    # Object read/write/delete
    if op in (
        "GetObject",
        "HeadObject",
        "GetObjectLockConfiguration",
        "GetObjectRetention",
        "GetObjectLegalHold",
    ):
        return "read_object"
    if op in (
        "PutObject",
        "CopyObject",
        "PutObjectLockConfiguration",
        "PutObjectRetention",
        "PutObjectLegalHold",
    ):
        return "write_object"
    if op in ("DeleteObject", "DeleteObjects"):
        return "delete_object"

    # Multipart uploads (object write flow)
    if op in (
        "CreateMultipartUpload",
        "UploadPart",
        "CompleteMultipartUpload",
        "AbortMultipartUpload",
    ):
        return "multipart_upload"

    # Bucket config surfaces
    if op in ("GetBucketVersioning", "PutBucketVersioning"):
        return "manage_bucket_versioning"
    if op in (
        "GetBucketEncryption",
        "PutBucketEncryption",
        "DeleteBucketEncryption",
    ):
        return "manage_bucket_encryption"
    if op in (
        "GetBucketLifecycleConfiguration",
        "PutBucketLifecycleConfiguration",
        "DeleteBucketLifecycle",
    ):
        return "manage_bucket_lifecycle"
    if op in ("GetBucketCors", "PutBucketCors", "DeleteBucketCors"):
        return "manage_bucket_cors"
    if op in ("GetBucketAcl", "PutBucketAcl"):
        return "manage_bucket_acl"
    if op in ("GetBucketTagging", "PutBucketTagging", "DeleteBucketTagging"):
        return "manage_bucket_tagging"
    if op in ("GetBucketPolicy", "PutBucketPolicy", "DeleteBucketPolicy"):
        return "manage_bucket_policy"
    if op in (
        "GetBucketNotificationConfiguration",
        "PutBucketNotificationConfiguration",
    ):
        return "manage_bucket_notifications"
    if op in ("GetBucketLogging", "PutBucketLogging"):
        return "manage_bucket_logging"
    if op in (
        "PutBucketReplication",
        "GetBucketReplication",
        "DeleteBucketReplication",
    ):
        return "manage_bucket_replication"

    return "misc"


def classify_dynamodb(op: str) -> str:
    if op in (
        "PutItem",
        "UpdateItem",
        "DeleteItem",
        "BatchWriteItem",
        "TransactWriteItems",
    ):
        return "write_item"
    if op in ("GetItem", "Query", "Scan", "BatchGetItem", "TransactGetItems"):
        return "read_item"

    if op in ("CreateTable", "UpdateTable", "DeleteTable"):
        return "manage_table"
    if op in ("DescribeTable", "ListTables", "DescribeEndpoints"):
        return "read_table"

    if op in ("DescribeTimeToLive", "UpdateTimeToLive"):
        return "manage_ttl"
    if op in ("DescribeContinuousBackups", "UpdateContinuousBackups"):
        return "manage_backups"

    if op in ("TagResource", "UntagResource", "ListTagsOfResource"):
        return "tagging"

    return "misc"


def classify_lambda(op: str) -> str:
    if op in ("Invoke",):
        return "invoke"

    if op in (
        "CreateFunction",
        "DeleteFunction",
        "UpdateFunctionCode",
        "UpdateFunctionConfiguration",
    ):
        return "manage_functions"
    if op in ("GetFunction", "ListFunctions", "ListVersionsByFunction"):
        return "read_functions"
    if op in ("PublishVersion",):
        return "versions"

    if op in ("AddPermission", "RemovePermission"):
        return "permissions"

    if op in (
        "CreateEventSourceMapping",
        "DeleteEventSourceMapping",
        "UpdateEventSourceMapping",
    ):
        return "manage_event_source_mappings"
    if op in ("GetEventSourceMapping", "ListEventSourceMappings"):
        return "read_event_source_mappings"

    if op in (
        "CreateFunctionUrlConfig",
        "UpdateFunctionUrlConfig",
        "DeleteFunctionUrlConfig",
    ):
        return "manage_function_urls"
    if op in ("GetFunctionUrlConfig", "ListFunctionUrlConfigs"):
        return "read_function_urls"

    if op in ("PublishLayerVersion", "DeleteLayerVersion"):
        return "manage_layers"
    if op in ("GetLayerVersion", "ListLayerVersions", "ListLayers"):
        return "read_layers"

    if op in ("TagResource", "UntagResource", "ListTags"):
        return "tagging"

    return "misc"


def classify_sqs(op: str) -> str:
    if op in ("SendMessage", "SendMessageBatch"):
        return "send_message"
    if op in ("ReceiveMessage",):
        return "receive_message"
    if op in ("DeleteMessage", "DeleteMessageBatch"):
        return "delete_message"
    if op in ("ChangeMessageVisibility", "ChangeMessageVisibilityBatch"):
        return "change_message_visibility"
    if op in (
        "CreateQueue",
        "DeleteQueue",
        "ListQueues",
        "GetQueueUrl",
        "GetQueueAttributes",
        "SetQueueAttributes",
        "PurgeQueue",
    ):
        return "manage_queues"
    if op in ("TagQueue", "UntagQueue", "ListQueueTags"):
        return "tagging"
    return "misc"


def classify_sns(op: str) -> str:
    if op in ("Publish",):
        return "publish"
    if op in ("Subscribe", "Unsubscribe", "ListSubscriptions"):
        return "manage_subscriptions"
    if op in ("CreateTopic", "DeleteTopic", "SetTopicAttributes"):
        return "manage_topics"
    if op in ("ListTopics", "GetTopicAttributes"):
        return "read_topics"
    if op in ("TagResource", "UntagResource", "ListTagsForResource"):
        return "tagging"
    return "misc"


def classify_cloudwatch(op: str) -> str:
    if op in ("PutMetricData",):
        return "write_metrics"
    if op in ("GetMetricStatistics", "GetMetricData", "ListMetrics"):
        return "read_metrics"
    if op in (
        "PutMetricAlarm",
        "PutCompositeAlarm",
        "DescribeAlarms",
        "DescribeAlarmsForMetric",
        "DescribeAlarmHistory",
        "DeleteAlarms",
        "SetAlarmState",
        "EnableAlarmActions",
        "DisableAlarmActions",
    ):
        return "manage_alarms"
    if op in (
        "PutDashboard",
        "GetDashboard",
        "DeleteDashboards",
        "ListDashboards",
    ):
        return "manage_dashboards"
    if op in ("TagResource", "UntagResource", "ListTagsForResource"):
        return "tagging"
    return "misc"


def classify_logs(op: str) -> str:
    if op in ("CreateLogGroup", "DeleteLogGroup", "DescribeLogGroups"):
        return "manage_log_groups"
    if op in ("CreateLogStream", "DeleteLogStream", "DescribeLogStreams"):
        return "manage_log_streams"
    if op in ("PutLogEvents",):
        return "write_events"
    if op in ("GetLogEvents", "FilterLogEvents"):
        return "read_events"
    if op in ("PutRetentionPolicy", "DeleteRetentionPolicy"):
        return "manage_retention"
    if op in (
        "PutSubscriptionFilter",
        "DeleteSubscriptionFilter",
        "DescribeSubscriptionFilters",
    ):
        return "manage_subscriptions"
    if op in ("PutMetricFilter", "DeleteMetricFilter", "DescribeMetricFilters"):
        return "manage_metric_filters"
    if op in ("PutDestination", "DeleteDestination", "DescribeDestinations"):
        return "manage_destinations"
    if op in ("StartQuery", "GetQueryResults", "StopQuery"):
        return "query"
    if op in (
        "TagLogGroup",
        "UntagLogGroup",
        "ListTagsLogGroup",
        "TagResource",
        "UntagResource",
        "ListTagsForResource",
    ):
        return "tagging"
    return "misc"


def classify_kinesis(op: str) -> str:
    if op in ("PutRecord", "PutRecords"):
        return "write_records"
    if op in ("GetShardIterator", "GetRecords"):
        return "read_records"
    if op.startswith("Describe") or op.startswith("List"):
        return "read_streams"
    if op in (
        "CreateStream",
        "DeleteStream",
        "MergeShards",
        "SplitShard",
        "UpdateShardCount",
    ):
        return "manage_streams"
    if op in ("StartStreamEncryption", "StopStreamEncryption"):
        return "manage_encryption"
    if op in ("EnableEnhancedMonitoring", "DisableEnhancedMonitoring"):
        return "manage_monitoring"
    if op in (
        "RegisterStreamConsumer",
        "DeregisterStreamConsumer",
        "ListStreamConsumers",
        "DescribeStreamConsumer",
    ):
        return "manage_consumers"
    if op in ("AddTagsToStream", "RemoveTagsFromStream", "ListTagsForStream"):
        return "tagging"
    return "misc"


def classify_ec2(op: str) -> str:
    # Explicit known buckets first
    if op.startswith("Describe"):
        return "describe"

    if op in ("CreateTags", "DeleteTags"):
        return "tagging"

    if "SecurityGroup" in op:
        return "manage_security_groups"

    if "KeyPair" in op:
        return "manage_keypairs"

    if any(
        k in op
        for k in (
            "Vpc",
            "Subnet",
            "Route",
            "Gateway",
            "InternetGateway",
            "NatGateway",
            "TransitGateway",
            "VpcEndpoint",
            "Network",
            "DhcpOptions",
            "Address",
            "EgressOnlyInternetGateway",
            "Vpn",
            "CarrierGateway",
            "CustomerGateway",
            "LocalGateway",
            "PrefixList",
        )
    ):
        return "manage_networking"

    if any(k in op for k in ("Volume", "Snapshot")):
        return "manage_storage"

    if "Image" in op:
        return "manage_images"

    if (
        "Instance" in op
        or op.endswith("Instances")
        or op
        in (
            "RunInstances",
            "StartInstances",
            "StopInstances",
            "RebootInstances",
            "TerminateInstances",
        )
    ):
        return "manage_instances"

    # Fallback to a generic-ish classifier to avoid unassigned operations.
    bucket = _generic_bucket(op)
    if bucket == "read":
        return "read"
    if bucket == "write":
        return "write"
    if bucket == "delete":
        return "delete"
    return "misc"


def classify_apigateway(op: str) -> str:
    if op.endswith("RestApi") or op.endswith("RestApis") or "RestApi" in op:
        if op.startswith("Get"):
            return "read_apis"
        return "manage_apis"

    if "Resource" in op:
        if op.startswith("Get"):
            return "read_resources"
        return "manage_resources"

    if "MethodResponse" in op or op.endswith("Method") or "Method" in op:
        if op.startswith("Get"):
            return "read_methods"
        return "manage_methods"

    if (
        "IntegrationResponse" in op
        or op.endswith("Integration")
        or "Integration" in op
    ):
        if op.startswith("Get"):
            return "read_integrations"
        return "manage_integrations"

    if "Deployment" in op:
        if op.startswith("Get"):
            return "read_deployments"
        return "manage_deployments"

    if op.endswith("Stage") or op.endswith("Stages") or "Stage" in op:
        if op.startswith("Get"):
            return "read_stages"
        return "manage_stages"

    if "Authorizer" in op:
        if op.startswith("Get"):
            return "read_authorizers"
        return "manage_authorizers"

    if "Model" in op:
        if op.startswith("Get"):
            return "read_models"
        return "manage_models"

    if "ApiKey" in op or "ApiKeys" in op:
        if op.startswith("Get"):
            return "read_api_keys"
        return "manage_api_keys"

    if "UsagePlanKey" in op or "UsagePlanKeys" in op:
        if op.startswith("Get"):
            return "read_usage_plan_keys"
        return "manage_usage_plan_keys"

    if "UsagePlan" in op:
        if op.startswith("Get"):
            return "read_usage_plans"
        return "manage_usage_plans"

    if "DomainName" in op or "DomainNames" in op:
        if op.startswith("Get"):
            return "read_domains"
        return "manage_domains"

    if "BasePathMapping" in op or "BasePathMappings" in op:
        if op.startswith("Get"):
            return "read_base_path_mappings"
        return "manage_base_path_mappings"

    if op in ("TagResource", "UntagResource", "GetTags"):
        return "tagging"

    return "misc"


def classify_apigatewayv2(op: str) -> str:
    if op.endswith("Api") or op.endswith("Apis") or "Api" in op:
        if op.startswith("Get"):
            return "read_apis"
        return "manage_apis"

    if "Route" in op:
        if op.startswith("Get"):
            return "read_routes"
        return "manage_routes"

    if "Integration" in op:
        if op.startswith("Get"):
            return "read_integrations"
        return "manage_integrations"

    if "Deployment" in op:
        if op.startswith("Get"):
            return "read_deployments"
        return "manage_deployments"

    if op.endswith("Stage") or op.endswith("Stages") or "Stage" in op:
        if op.startswith("Get"):
            return "read_stages"
        return "manage_stages"

    if "Authorizer" in op:
        if op.startswith("Get"):
            return "read_authorizers"
        return "manage_authorizers"

    if op in ("TagResource", "UntagResource", "GetTags"):
        return "tagging"

    return "misc"


CORE_CLASSIFIERS: dict[str, Callable[[str], str]] = {
    "s3": classify_s3,
    "ec2": classify_ec2,
    "sns": classify_sns,
    "sqs": classify_sqs,
    "lambda": classify_lambda,
    "apigateway": classify_apigateway,
    "apigatewayv2": classify_apigatewayv2,
    "cloudwatch": classify_cloudwatch,
    "logs": classify_logs,
    "kinesis": classify_kinesis,
    "dynamodb": classify_dynamodb,
}


def build_capabilities() -> tuple[
    dict[str, dict[str, object]], dict[str, dict[str, list[str]]]
]:
    """
    Returns:
      - CAPABILITY_TO_OPERATIONS mapping (capability -> {service, operation_sets})
      - per_service_groups (service -> suffix -> [ops]) for catalog rendering
    """
    per_service: dict[str, DefaultDict[str, list[str]]] = {}
    for svc, ops in AWS_SERVICE_OPERATIONS.items():
        groups: DefaultDict[str, list[str]] = defaultdict(list)
        classifier = CORE_CLASSIFIERS.get(svc)
        for op in ops:
            if classifier is not None:
                suffix = classifier(op)
            else:
                suffix = _generic_bucket(op)
            groups[suffix].append(op)
        per_service[svc] = groups

    # Validate core services coverage (no op left unbucketed)
    missing: list[str] = []
    for svc in sorted(CORE_SERVICES):
        if svc not in per_service:
            missing.append(f"{svc}: missing from AWS_SERVICE_OPERATIONS")
            continue
        assigned = {op for ops in per_service[svc].values() for op in ops}
        for op in AWS_SERVICE_OPERATIONS.get(svc, []):
            if op not in assigned:
                missing.append(f"{svc}.{op}: unassigned")
    if missing:
        raise SystemExit(
            "Capability derivation failed (core coverage):\n"
            + "\n".join("  - " + m for m in missing)
        )

    cap_to_ops: dict[str, dict[str, object]] = {}
    for svc, groups in per_service.items():
        for suffix, ops in groups.items():
            cap = f"{svc}.{suffix}"
            cap_to_ops[cap] = {
                "service": svc,
                "operation_sets": [[op] for op in sorted(set(ops))],
            }

    # Stable ordering by key (writer also sorts)
    cap_to_ops = dict(sorted(cap_to_ops.items(), key=lambda kv: kv[0]))

    # Convert per_service groups to plain dicts with sorted ops
    per_service_sorted: dict[str, dict[str, list[str]]] = {}
    for svc, groups in per_service.items():
        per_service_sorted[svc] = {
            suffix: sorted(set(ops)) for suffix, ops in sorted(groups.items())
        }
    return cap_to_ops, per_service_sorted


def write_capabilities_py(cap_to_ops: dict[str, dict[str, object]]) -> Path:
    out = AWS_DIR / "capabilities_data.py"
    rendered = pprint.pformat(cap_to_ops, sort_dicts=True, width=120)
    core_literal = "{%s}" % ", ".join(repr(s) for s in sorted(CORE_SERVICES))
    text = "\n".join(
        [
            "from __future__ import annotations",
            "",
            '"""',
            "AUTO-GENERATED FILE. DO NOT EDIT BY HAND.",
            "",
            "Generated by: cloudenv/scripts/derive_capability_to_operations.py",
            "Source: cloudenv/server/providers/aws/operations_data.py",
            '"""',
            "",
            f"CORE_SERVICES: set[str] = {core_literal}",
            "",
            "CAPABILITY_TO_OPERATIONS: dict[str, dict[str, object]] = "
            + rendered,
            "",
        ]
    )
    out.write_text(text, encoding="utf-8")
    return out


def write_catalog_md(per_service: dict[str, dict[str, list[str]]]) -> Path:
    out = AWS_DIR / "capabilities_catalog.md"
    lines: list[str] = [
        "# Capability Catalog",
        "",
        "_Auto-generated by `cloudenv/scripts/derive_capability_to_operations.py`._",
        "",
        "Each capability maps to one or more **operation sets**; this generator emits one operation per set (OR semantics).",
        "",
    ]
    for svc in sorted(per_service.keys()):
        lines.append(f"## {svc}")
        lines.append("")
        groups = per_service[svc]
        for suffix in sorted(groups.keys()):
            ops = groups[suffix]
            cap = f"{svc}.{suffix}"
            lines.append(f"### `{cap}` ({len(ops)} ops)")
            lines.append("")
            for op in ops:
                lines.append(f"- `{op}`")
            lines.append("")
    out.write_text("\n".join(lines).rstrip() + "\n", encoding="utf-8")
    return out


def main() -> None:
    cap_to_ops, per_service = build_capabilities()
    py_path = write_capabilities_py(cap_to_ops)
    md_path = write_catalog_md(per_service)
    print(f"Wrote {py_path.relative_to(ROOT)}")
    print(f"Wrote {md_path.relative_to(ROOT)}")
    print(f"Total capabilities: {len(cap_to_ops)}")


if __name__ == "__main__":
    main()