AWSTemplateFormatVersion: '2010-09-09'
Description: Amazon MWAA Environment Cost Review Remediator - review-only inventory across enabled Regions.

Parameters:
  ScheduleEnabled:
    Type: String
    Default: 'false'
    AllowedValues: ['true', 'false']
  ScheduleExpression:
    Type: String
    Default: cron(0 13 ? * MON *)
  ExcludedRegionsCsv:
    Type: String
    Default: ap-northeast-3
  IncludeEnvironmentDetails:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
  IncludeTags:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
  MaxEnvironmentsPerRegion:
    Type: Number
    Default: 25
    MinValue: 1
    MaxValue: 200

Conditions:
  CreateSchedule: !Equals [!Ref ScheduleEnabled, 'true']

Resources:
  MwaaCostReviewRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: [lambda.amazonaws.com]
            Action: [sts:AssumeRole]
      Policies:
        - PolicyName: mwaa-cost-review-policy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Sid: LambdaLogging
                Effect: Allow
                Action: [logs:CreateLogGroup, logs:CreateLogStream, logs:PutLogEvents]
                Resource: '*'
              - Sid: RegionDiscovery
                Effect: Allow
                Action: [ec2:DescribeRegions]
                Resource: '*'
              - Sid: MwaaReadOnly
                Effect: Allow
                Action:
                  - airflow:ListEnvironments
                  - airflow:GetEnvironment
                  - airflow:ListTagsForResource
                Resource: '*'

  MwaaCostReviewFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: !Sub mwaa-cost-review-${AWS::StackName}
      Runtime: python3.12
      Handler: index.lambda_handler
      Role: !GetAtt MwaaCostReviewRole.Arn
      Timeout: 900
      MemorySize: 256
      Environment:
        Variables:
          EXCLUDED_REGIONS_CSV: !Ref ExcludedRegionsCsv
          INCLUDE_ENVIRONMENT_DETAILS: !Ref IncludeEnvironmentDetails
          INCLUDE_TAGS: !Ref IncludeTags
          MAX_ENVIRONMENTS_PER_REGION: !Ref MaxEnvironmentsPerRegion
      Code:
        ZipFile: |
          import boto3
          import os
          import traceback
          from datetime import datetime, timezone
          from botocore.exceptions import ClientError, EndpointConnectionError

          def now():
              return datetime.now(timezone.utc).isoformat()

          def as_bool(value, default=False):
              if value is None:
                  return default
              return str(value).strip().lower() in ("1", "true", "yes", "y")

          def as_int(value, default):
              try:
                  return int(value)
              except Exception:
                  return default

          def csv_set(value):
              return set(x.strip() for x in str(value or "").split(",") if x.strip())

          def enabled_regions():
              ec2 = boto3.client("ec2")
              response = ec2.describe_regions(AllRegions=True)
              out = []
              for r in response.get("Regions", []):
                  status = r.get("OptInStatus")
                  if status in (None, "opt-in-not-required", "opted-in"):
                      out.append(r["RegionName"])
              return sorted(out)

          def list_environments(client, max_envs):
              envs = []
              token = None
              while True:
                  # MWAA ListEnvironments has a lower MaxResults ceiling than many AWS APIs.
                  # Use 25 per API call and paginate until MaxEnvironmentsPerRegion is reached.
                  kwargs = {"MaxResults": min(max_envs, 25)}
                  if token:
                      kwargs["NextToken"] = token
                  response = client.list_environments(**kwargs)
                  envs.extend(response.get("Environments", []) or [])
                  if len(envs) >= max_envs:
                      return envs[:max_envs], True
                  token = response.get("NextToken")
                  if not token:
                      return envs, False

          def logging_summary(logging_cfg):
              logging_cfg = logging_cfg if isinstance(logging_cfg, dict) else {}
              result = {}
              enabled_count = 0
              verbose_count = 0
              for key in ["DagProcessingLogs", "SchedulerLogs", "TaskLogs", "WebserverLogs", "WorkerLogs"]:
                  cfg = logging_cfg.get(key, {}) if isinstance(logging_cfg.get(key), dict) else {}
                  enabled = cfg.get("Enabled")
                  level = cfg.get("LogLevel")
                  result[key] = {"enabled": enabled, "log_level": level}
                  if enabled is True:
                      enabled_count += 1
                      if str(level or "").upper() in ("DEBUG", "INFO", "WARNING"):
                          verbose_count += 1
              return result, enabled_count, verbose_count

          def summarize_environment(region, client, name, include_tags):
              item = {"region": region, "name": name, "review_flags": []}
              try:
                  env = client.get_environment(Name=name).get("Environment", {}) or {}
                  logging, enabled_logs, verbose_logs = logging_summary(env.get("LoggingConfiguration"))
                  item.update({
                      "arn": env.get("Arn"),
                      "status": env.get("Status"),
                      "airflow_version": env.get("AirflowVersion"),
                      "environment_class": env.get("EnvironmentClass"),
                      "created_at": str(env.get("CreatedAt")) if env.get("CreatedAt") else None,
                      "last_update_status": (env.get("LastUpdate") or {}).get("Status"),
                      "webserver_access_mode": env.get("WebserverAccessMode"),
                      "min_workers": env.get("MinWorkers"),
                      "max_workers": env.get("MaxWorkers"),
                      "schedulers": env.get("Schedulers"),
                      "weekly_maintenance_window_start": env.get("WeeklyMaintenanceWindowStart"),
                      "source_bucket_arn_present": bool(env.get("SourceBucketArn")),
                      "dag_s3_path": env.get("DagS3Path"),
                      "plugins_s3_path_present": bool(env.get("PluginsS3Path")),
                      "requirements_s3_path_present": bool(env.get("RequirementsS3Path")),
                      "startup_script_s3_path_present": bool(env.get("StartupScriptS3Path")),
                      "kms_key_present": bool(env.get("KmsKey")),
                      "execution_role_arn_present": bool(env.get("ExecutionRoleArn")),
                      "security_group_count": len(((env.get("NetworkConfiguration") or {}).get("SecurityGroupIds") or [])),
                      "subnet_count": len(((env.get("NetworkConfiguration") or {}).get("SubnetIds") or [])),
                      "logging": logging,
                      "logging_enabled_categories": enabled_logs,
                      "verbose_logging_categories": verbose_logs
                  })

                  arn = env.get("Arn")
                  if include_tags and arn:
                      try:
                          tags = client.list_tags_for_resource(ResourceArn=arn).get("Tags", {})
                          item["tags"] = tags
                          item["tag_count"] = len(tags)
                      except ClientError as e:
                          item["tags_error"] = e.response.get("Error", {}).get("Code", "ClientError")

              except ClientError as e:
                  item["get_error"] = e.response.get("Error", {}).get("Code", "ClientError")
              except Exception as e:
                  item["get_error"] = str(e)

              flags = []
              status = str(item.get("status") or "").upper()
              env_class = str(item.get("environment_class") or "")
              web_mode = str(item.get("webserver_access_mode") or "")
              max_workers = item.get("max_workers") or 0
              min_workers = item.get("min_workers") or 0
              schedulers = item.get("schedulers") or 0

              if status == "AVAILABLE":
                  flags.append("available_environment_review")
              elif status:
                  flags.append("environment_state_" + status.lower())
              if env_class and env_class not in ("mw1.small", "mw1.micro"):
                  flags.append("larger_environment_class_review")
              if max_workers and max_workers >= 5:
                  flags.append("high_max_workers_review")
              if min_workers and min_workers >= 2:
                  flags.append("multiple_min_workers_review")
              if schedulers and schedulers >= 2:
                  flags.append("multiple_schedulers_review")
              if item.get("logging_enabled_categories", 0) > 0:
                  flags.append("logging_enabled_review")
              if item.get("verbose_logging_categories", 0) > 0:
                  flags.append("verbose_logging_review")
              if web_mode == "PUBLIC_ONLY":
                  flags.append("public_webserver_review")
              if item.get("last_update_status") and str(item.get("last_update_status")).upper() not in ("SUCCESS",):
                  flags.append("last_update_not_success_review")

              item["review_flags"] = sorted(set(flags))
              return item

          def lambda_handler(event, context):
              event = event if isinstance(event, dict) else {}
              excluded = csv_set(event.get("excluded_regions_csv", os.environ.get("EXCLUDED_REGIONS_CSV", "ap-northeast-3")))
              include_details = as_bool(event.get("include_details", os.environ.get("INCLUDE_ENVIRONMENT_DETAILS", "true")), True)
              include_tags = as_bool(event.get("include_tags", os.environ.get("INCLUDE_TAGS", "true")), True)
              max_envs = as_int(event.get("max_environments_per_region", os.environ.get("MAX_ENVIRONMENTS_PER_REGION", "25")), 25)

              result = {
                  "module": "amazon_mwaa_environment_cost_review",
                  "review_only": True,
                  "started_at": now(),
                  "ended_at": None,
                  "excluded_regions": sorted(excluded),
                  "totals": {
                      "regions_scanned": 0,
                      "regions_skipped": 0,
                      "environments": 0,
                      "available_environments": 0,
                      "non_available_environments": 0,
                      "larger_environment_class": 0,
                      "public_webserver": 0,
                      "logging_enabled_environments": 0,
                      "verbose_logging_environments": 0,
                      "total_min_workers": 0,
                      "total_max_workers": 0,
                      "total_schedulers": 0,
                      "regions_with_environments": 0,
                      "regions_with_review_flags": 0,
                      "truncated_regions": 0,
                      "access_denied": 0,
                      "failed": 0
                  },
                  "regions": [],
                  "environments": []
              }

              try:
                  regions = enabled_regions()
              except Exception as e:
                  result["fatal_error"] = str(e)
                  result["traceback"] = traceback.format_exc()
                  result["ended_at"] = now()
                  return result

              for region in regions:
                  if region in excluded:
                      result["totals"]["regions_skipped"] += 1
                      result["regions"].append({"region": region, "skipped": True, "reason": "excluded"})
                      continue

                  rec = {
                      "region": region,
                      "skipped": False,
                      "environments": 0,
                      "available_environments": 0,
                      "non_available_environments": 0,
                      "larger_environment_class": 0,
                      "public_webserver": 0,
                      "logging_enabled_environments": 0,
                      "verbose_logging_environments": 0,
                      "min_workers": 0,
                      "max_workers": 0,
                      "schedulers": 0,
                      "review_flags": [],
                      "truncated": False,
                      "error": None
                  }

                  try:
                      client = boto3.client("mwaa", region_name=region)
                      names, truncated = list_environments(client, max_envs)
                      rec["truncated"] = truncated
                      if truncated:
                          result["totals"]["truncated_regions"] += 1

                      details = []
                      for name in names:
                          d = summarize_environment(region, client, name, include_tags)
                          details.append(d)
                          rec["environments"] += 1
                          result["totals"]["environments"] += 1

                          if str(d.get("status") or "").upper() == "AVAILABLE":
                              rec["available_environments"] += 1
                              result["totals"]["available_environments"] += 1
                          else:
                              rec["non_available_environments"] += 1
                              result["totals"]["non_available_environments"] += 1

                          env_class = str(d.get("environment_class") or "")
                          if env_class and env_class not in ("mw1.small", "mw1.micro"):
                              rec["larger_environment_class"] += 1
                              result["totals"]["larger_environment_class"] += 1
                          if d.get("webserver_access_mode") == "PUBLIC_ONLY":
                              rec["public_webserver"] += 1
                              result["totals"]["public_webserver"] += 1
                          if (d.get("logging_enabled_categories") or 0) > 0:
                              rec["logging_enabled_environments"] += 1
                              result["totals"]["logging_enabled_environments"] += 1
                          if (d.get("verbose_logging_categories") or 0) > 0:
                              rec["verbose_logging_environments"] += 1
                              result["totals"]["verbose_logging_environments"] += 1

                          for source_key, total_key in [
                              ("min_workers", "total_min_workers"),
                              ("max_workers", "total_max_workers"),
                              ("schedulers", "total_schedulers")
                          ]:
                              value = d.get(source_key) or 0
                              rec[source_key] += value
                              result["totals"][total_key] += value

                          rec["review_flags"].extend(d.get("review_flags", []))

                      rec["review_flags"] = sorted(set(rec["review_flags"]))

                      if rec["environments"] > 0:
                          result["totals"]["regions_with_environments"] += 1
                      if rec["review_flags"]:
                          result["totals"]["regions_with_review_flags"] += 1

                      if include_details:
                          result["environments"].extend(details)

                      result["totals"]["regions_scanned"] += 1
                      result["regions"].append(rec)

                  except ClientError as e:
                      code = e.response.get("Error", {}).get("Code", "ClientError")
                      rec["error"] = code
                      if code in ("AccessDenied", "AccessDeniedException", "UnauthorizedOperation"):
                          result["totals"]["access_denied"] += 1
                      else:
                          result["totals"]["failed"] += 1
                      result["regions"].append(rec)

                  except EndpointConnectionError:
                      rec["error"] = "EndpointConnectionError"
                      result["totals"]["failed"] += 1
                      result["regions"].append(rec)

                  except Exception as e:
                      rec["error"] = str(e)
                      rec["traceback"] = traceback.format_exc()
                      result["totals"]["failed"] += 1
                      result["regions"].append(rec)

              result["ended_at"] = now()
              return result

  MwaaCostReviewSchedule:
    Type: AWS::Events::Rule
    Condition: CreateSchedule
    Properties:
      Description: Scheduled review-only Amazon MWAA environment cost review.
      ScheduleExpression: !Ref ScheduleExpression
      State: ENABLED
      Targets:
        - Arn: !GetAtt MwaaCostReviewFunction.Arn
          Id: MwaaCostReviewFunctionTarget
          Input: '{"scheduled":true}'

  MwaaCostReviewSchedulePermission:
    Type: AWS::Lambda::Permission
    Condition: CreateSchedule
    Properties:
      FunctionName: !Ref MwaaCostReviewFunction
      Action: lambda:InvokeFunction
      Principal: events.amazonaws.com
      SourceArn: !GetAtt MwaaCostReviewSchedule.Arn

Outputs:
  ReviewFunctionName:
    Description: Lambda function name.
    Value: !Ref MwaaCostReviewFunction
  ReviewFunctionArn:
    Description: Lambda function ARN.
    Value: !GetAtt MwaaCostReviewFunction.Arn
  ReviewRoleArn:
    Description: IAM role ARN used by the review Lambda.
    Value: !GetAtt MwaaCostReviewRole.Arn
