AWSTemplateFormatVersion: '2010-09-09'
Description: Amazon EMR Cluster 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
  IncludeClusterDetails:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
  IncludeInstanceGroups:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
  IncludeInstanceFleets:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
  IncludeInstances:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
  MaxClustersPerRegion:
    Type: Number
    Default: 200
    MinValue: 1
    MaxValue: 500
  MaxInstancesPerCluster:
    Type: Number
    Default: 200
    MinValue: 1
    MaxValue: 500
  LongRunningDaysThreshold:
    Type: Number
    Default: 7
    MinValue: 1
    MaxValue: 365

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

Resources:
  EmrCostReviewRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: [lambda.amazonaws.com]
            Action: [sts:AssumeRole]
      Policies:
        - PolicyName: emr-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: EmrReadOnly
                Effect: Allow
                Action:
                  - elasticmapreduce:ListClusters
                  - elasticmapreduce:DescribeCluster
                  - elasticmapreduce:ListInstanceGroups
                  - elasticmapreduce:ListInstanceFleets
                  - elasticmapreduce:ListInstances
                  - elasticmapreduce:GetAutoTerminationPolicy
                Resource: '*'

  EmrCostReviewFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: !Sub emr-cost-review-${AWS::StackName}
      Runtime: python3.12
      Handler: index.lambda_handler
      Role: !GetAtt EmrCostReviewRole.Arn
      Timeout: 900
      MemorySize: 256
      Environment:
        Variables:
          EXCLUDED_REGIONS_CSV: !Ref ExcludedRegionsCsv
          INCLUDE_CLUSTER_DETAILS: !Ref IncludeClusterDetails
          INCLUDE_INSTANCE_GROUPS: !Ref IncludeInstanceGroups
          INCLUDE_INSTANCE_FLEETS: !Ref IncludeInstanceFleets
          INCLUDE_INSTANCES: !Ref IncludeInstances
          MAX_CLUSTERS_PER_REGION: !Ref MaxClustersPerRegion
          MAX_INSTANCES_PER_CLUSTER: !Ref MaxInstancesPerCluster
          LONG_RUNNING_DAYS_THRESHOLD: !Ref LongRunningDaysThreshold
      Code:
        ZipFile: |
          import boto3
          import os
          import traceback
          from datetime import datetime, timezone
          from botocore.exceptions import ClientError, EndpointConnectionError

          ACTIVE_STATES = ["STARTING", "BOOTSTRAPPING", "RUNNING", "WAITING", "TERMINATING"]

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

          def now_iso():
              return now().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 call_with_marker(method, result_key, max_items, **kwargs):
              items = []
              marker = None
              while True:
                  call_args = dict(kwargs)
                  if marker:
                      call_args["Marker"] = marker
                  response = method(**call_args)
                  items.extend(response.get(result_key, []) or [])
                  if len(items) >= max_items:
                      return items[:max_items], True
                  marker = response.get("Marker")
                  if not marker:
                      return items, False

          def list_clusters(emr, max_clusters):
              clusters = []
              marker = None
              while True:
                  kwargs = {"ClusterStates": ACTIVE_STATES}
                  if marker:
                      kwargs["Marker"] = marker
                  response = emr.list_clusters(**kwargs)
                  clusters.extend(response.get("Clusters", []) or [])
                  if len(clusters) >= max_clusters:
                      return clusters[:max_clusters], True
                  marker = response.get("Marker")
                  if not marker:
                      return clusters, False

          def age_days_from_datetime(value):
              if not value:
                  return None
              try:
                  if value.tzinfo is None:
                      value = value.replace(tzinfo=timezone.utc)
                  return round((now() - value).total_seconds() / 86400, 2)
              except Exception:
                  return None

          def get_auto_termination(emr, cluster_id):
              try:
                  response = emr.get_auto_termination_policy(ClusterId=cluster_id)
                  policy = response.get("AutoTerminationPolicy", {})
                  return {
                      "present": bool(policy),
                      "idle_timeout_seconds": policy.get("IdleTimeout")
                  }
              except ClientError as e:
                  code = e.response.get("Error", {}).get("Code", "ClientError")
                  if code in ("ValidationException", "InvalidRequestException"):
                      return {"present": False, "idle_timeout_seconds": None}
                  return {"error": code}
              except Exception as e:
                  return {"error": str(e)}

          def summarize_cluster(region, emr, cluster_summary, include_groups, include_fleets, include_instances, max_instances, long_days):
              cluster_id = cluster_summary.get("Id")
              status = cluster_summary.get("Status", {}) or {}
              timeline = status.get("Timeline", {}) or {}
              creation = timeline.get("CreationDateTime")
              state = status.get("State")

              item = {
                  "region": region,
                  "cluster_id": cluster_id,
                  "name": cluster_summary.get("Name"),
                  "state": state,
                  "normalized_instance_hours": cluster_summary.get("NormalizedInstanceHours"),
                  "creation_time": str(creation) if creation else None,
                  "age_days": age_days_from_datetime(creation),
                  "review_flags": []
              }

              try:
                  detail = emr.describe_cluster(ClusterId=cluster_id).get("Cluster", {})
                  status_d = detail.get("Status", {}) or {}
                  timeline_d = status_d.get("Timeline", {}) or {}
                  creation_d = timeline_d.get("CreationDateTime") or creation

                  item.update({
                      "arn": detail.get("ClusterArn"),
                      "release_label": detail.get("ReleaseLabel"),
                      "auto_terminate": detail.get("AutoTerminate"),
                      "termination_protected": detail.get("TerminationProtected"),
                      "visible_to_all_users": detail.get("VisibleToAllUsers"),
                      "log_uri_present": bool(detail.get("LogUri")),
                      "applications": [a.get("Name") for a in detail.get("Applications", []) or []],
                      "applications_count": len(detail.get("Applications", []) or []),
                      "tags": {t.get("Key"): t.get("Value") for t in detail.get("Tags", []) or [] if t.get("Key")},
                      "tag_count": len(detail.get("Tags", []) or []),
                      "service_role": detail.get("ServiceRole"),
                      "job_flow_role": detail.get("JobFlowRole"),
                      "scale_down_behavior": detail.get("ScaleDownBehavior"),
                      "ebs_root_volume_size": detail.get("EbsRootVolumeSize"),
                      "step_concurrency_level": detail.get("StepConcurrencyLevel"),
                      "managed_scaling_policy_present": bool(detail.get("ManagedScalingPolicy")),
                      "placement_availability_zone": (detail.get("Ec2InstanceAttributes") or {}).get("Ec2AvailabilityZone"),
                      "requested_ami_version": detail.get("RequestedAmiVersion"),
                      "running_ami_version": detail.get("RunningAmiVersion")
                  })

                  item["creation_time"] = str(creation_d) if creation_d else item.get("creation_time")
                  item["age_days"] = age_days_from_datetime(creation_d)

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

              item["auto_termination_policy"] = get_auto_termination(emr, cluster_id)

              if include_groups:
                  try:
                      groups, trunc = call_with_marker(emr.list_instance_groups, "InstanceGroups", 100, ClusterId=cluster_id)
                      item["instance_groups_count"] = len(groups)
                      item["instance_groups_truncated"] = trunc
                      item["instance_groups"] = [
                          {
                              "id": g.get("Id"),
                              "name": g.get("Name"),
                              "type": g.get("InstanceGroupType"),
                              "instance_type": g.get("InstanceType"),
                              "requested_count": g.get("RequestedInstanceCount"),
                              "running_count": g.get("RunningInstanceCount"),
                              "market": g.get("Market"),
                              "state": (g.get("Status") or {}).get("State")
                          }
                          for g in groups
                      ]
                      item["requested_instance_count"] = sum(g.get("RequestedInstanceCount") or 0 for g in groups)
                      item["running_instance_count"] = sum(g.get("RunningInstanceCount") or 0 for g in groups)
                  except ClientError as e:
                      item["instance_groups_error"] = e.response.get("Error", {}).get("Code", "ClientError")
                  except Exception as e:
                      item["instance_groups_error"] = str(e)

              if include_fleets:
                  try:
                      fleets, trunc = call_with_marker(emr.list_instance_fleets, "InstanceFleets", 100, ClusterId=cluster_id)
                      item["instance_fleets_count"] = len(fleets)
                      item["instance_fleets_truncated"] = trunc
                      item["instance_fleets"] = [
                          {
                              "id": f.get("Id"),
                              "name": f.get("Name"),
                              "type": f.get("InstanceFleetType"),
                              "target_on_demand": (f.get("TargetOnDemandCapacity") or 0),
                              "target_spot": (f.get("TargetSpotCapacity") or 0),
                              "provisioned_on_demand": (f.get("ProvisionedOnDemandCapacity") or 0),
                              "provisioned_spot": (f.get("ProvisionedSpotCapacity") or 0),
                              "state": (f.get("Status") or {}).get("State")
                          }
                          for f in fleets
                      ]
                      item["fleet_target_capacity"] = sum((f.get("TargetOnDemandCapacity") or 0) + (f.get("TargetSpotCapacity") or 0) for f in fleets)
                      item["fleet_provisioned_capacity"] = sum((f.get("ProvisionedOnDemandCapacity") or 0) + (f.get("ProvisionedSpotCapacity") or 0) for f in fleets)
                  except ClientError as e:
                      item["instance_fleets_error"] = e.response.get("Error", {}).get("Code", "ClientError")
                  except Exception as e:
                      item["instance_fleets_error"] = str(e)

              if include_instances:
                  try:
                      instances, trunc = call_with_marker(emr.list_instances, "Instances", max_instances, ClusterId=cluster_id)
                      item["sampled_instances"] = len(instances)
                      item["instances_truncated"] = trunc
                      item["sampled_instance_types"] = sorted(set([i.get("InstanceType") for i in instances if i.get("InstanceType")]))
                  except ClientError as e:
                      item["instances_error"] = e.response.get("Error", {}).get("Code", "ClientError")
                  except Exception as e:
                      item["instances_error"] = str(e)

              flags = []
              state_upper = str(item.get("state") or "").upper()
              if state_upper in ("RUNNING", "WAITING"):
                  flags.append("active_or_waiting_cluster_review")
              if state_upper == "WAITING":
                  flags.append("waiting_cluster_review")
              if item.get("age_days") is not None and item.get("age_days") >= long_days:
                  flags.append("long_running_cluster_review")
              if item.get("auto_termination_policy", {}).get("present") is not True:
                  flags.append("missing_auto_termination_policy_review")
              if item.get("termination_protected") is True:
                  flags.append("termination_protected_review")
              if (item.get("running_instance_count") or 0) > 0:
                  flags.append("running_instances_review")
              if (item.get("fleet_provisioned_capacity") or 0) > 0:
                  flags.append("provisioned_fleet_capacity_review")
              if item.get("managed_scaling_policy_present") is True:
                  flags.append("managed_scaling_policy_review")
              if item.get("normalized_instance_hours") and item.get("normalized_instance_hours") > 0:
                  flags.append("normalized_instance_hours_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_CLUSTER_DETAILS", "true")), True)
              include_groups = as_bool(event.get("include_instance_groups", os.environ.get("INCLUDE_INSTANCE_GROUPS", "true")), True)
              include_fleets = as_bool(event.get("include_instance_fleets", os.environ.get("INCLUDE_INSTANCE_FLEETS", "true")), True)
              include_instances = as_bool(event.get("include_instances", os.environ.get("INCLUDE_INSTANCES", "true")), True)
              max_clusters = as_int(event.get("max_clusters_per_region", os.environ.get("MAX_CLUSTERS_PER_REGION", "200")), 200)
              max_instances = as_int(event.get("max_instances_per_cluster", os.environ.get("MAX_INSTANCES_PER_CLUSTER", "200")), 200)
              long_days = as_int(event.get("long_running_days_threshold", os.environ.get("LONG_RUNNING_DAYS_THRESHOLD", "7")), 7)

              result = {
                  "module": "amazon_emr_cluster_cost_review",
                  "review_only": True,
                  "started_at": now_iso(),
                  "ended_at": None,
                  "excluded_regions": sorted(excluded),
                  "active_states_scanned": ACTIVE_STATES,
                  "totals": {
                      "regions_scanned": 0,
                      "regions_skipped": 0,
                      "clusters": 0,
                      "running_clusters": 0,
                      "waiting_clusters": 0,
                      "long_running_clusters": 0,
                      "clusters_without_auto_termination_policy": 0,
                      "termination_protected_clusters": 0,
                      "requested_instance_count": 0,
                      "running_instance_count": 0,
                      "fleet_target_capacity": 0,
                      "fleet_provisioned_capacity": 0,
                      "sampled_instances": 0,
                      "regions_with_clusters": 0,
                      "regions_with_review_flags": 0,
                      "truncated_regions": 0,
                      "access_denied": 0,
                      "failed": 0
                  },
                  "regions": [],
                  "clusters": []
              }

              try:
                  regions = enabled_regions()
              except Exception as e:
                  result["fatal_error"] = str(e)
                  result["traceback"] = traceback.format_exc()
                  result["ended_at"] = now_iso()
                  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,
                      "clusters": 0,
                      "running_clusters": 0,
                      "waiting_clusters": 0,
                      "long_running_clusters": 0,
                      "clusters_without_auto_termination_policy": 0,
                      "termination_protected_clusters": 0,
                      "requested_instance_count": 0,
                      "running_instance_count": 0,
                      "fleet_target_capacity": 0,
                      "fleet_provisioned_capacity": 0,
                      "sampled_instances": 0,
                      "review_flags": [],
                      "truncated": False,
                      "error": None
                  }

                  try:
                      emr = boto3.client("emr", region_name=region)
                      clusters, truncated = list_clusters(emr, max_clusters)
                      rec["truncated"] = truncated
                      if truncated:
                          result["totals"]["truncated_regions"] += 1

                      details = []
                      for c in clusters:
                          d = summarize_cluster(region, emr, c, include_groups, include_fleets, include_instances, max_instances, long_days)
                          details.append(d)

                          rec["clusters"] += 1
                          result["totals"]["clusters"] += 1

                          state = str(d.get("state") or "").upper()
                          if state == "RUNNING":
                              rec["running_clusters"] += 1
                              result["totals"]["running_clusters"] += 1
                          if state == "WAITING":
                              rec["waiting_clusters"] += 1
                              result["totals"]["waiting_clusters"] += 1
                          if d.get("age_days") is not None and d.get("age_days") >= long_days:
                              rec["long_running_clusters"] += 1
                              result["totals"]["long_running_clusters"] += 1
                          if d.get("auto_termination_policy", {}).get("present") is not True:
                              rec["clusters_without_auto_termination_policy"] += 1
                              result["totals"]["clusters_without_auto_termination_policy"] += 1
                          if d.get("termination_protected") is True:
                              rec["termination_protected_clusters"] += 1
                              result["totals"]["termination_protected_clusters"] += 1

                          for k in ("requested_instance_count", "running_instance_count", "fleet_target_capacity", "fleet_provisioned_capacity", "sampled_instances"):
                              value = d.get(k) or 0
                              rec[k] += value
                              result["totals"][k] += value

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

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

                      if rec["clusters"] > 0:
                          result["totals"]["regions_with_clusters"] += 1
                      if rec["review_flags"]:
                          result["totals"]["regions_with_review_flags"] += 1

                      if include_details:
                          result["clusters"].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_iso()
              return result

  EmrCostReviewSchedule:
    Type: AWS::Events::Rule
    Condition: CreateSchedule
    Properties:
      Description: Scheduled review-only Amazon EMR cluster cost review.
      ScheduleExpression: !Ref ScheduleExpression
      State: ENABLED
      Targets:
        - Arn: !GetAtt EmrCostReviewFunction.Arn
          Id: EmrCostReviewFunctionTarget
          Input: '{"scheduled":true}'

  EmrCostReviewSchedulePermission:
    Type: AWS::Lambda::Permission
    Condition: CreateSchedule
    Properties:
      FunctionName: !Ref EmrCostReviewFunction
      Action: lambda:InvokeFunction
      Principal: events.amazonaws.com
      SourceArn: !GetAtt EmrCostReviewSchedule.Arn

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