AWSTemplateFormatVersion: '2010-09-09'
Description: >
  EMR / EMR Serverless / EMR on EKS Cost Review Remediator.
  Review-only inventory of EMR clusters, EMR Serverless applications, and EMR on EKS virtual clusters.

Parameters:
  ProjectName:
    Type: String
    Default: emr-cost-review
  IncludeDetails:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
  MaxClustersPerRegion:
    Type: Number
    Default: 200
  MaxStepsPerCluster:
    Type: Number
    Default: 50
  ExcludedRegionsCsv:
    Type: String
    Default: ap-northeast-3
  ScheduleEnabled:
    Type: String
    Default: 'false'
    AllowedValues: ['true', 'false']
  ScheduleExpression:
    Type: String
    Default: rate(30 days)

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

Resources:
  ReviewRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal: { Service: lambda.amazonaws.com }
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyName: EmrCostReviewReadOnlyPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action: ec2:DescribeRegions
                Resource: '*'
              - Effect: Allow
                Action:
                  - elasticmapreduce:ListClusters
                  - elasticmapreduce:DescribeCluster
                  - elasticmapreduce:ListInstanceGroups
                  - elasticmapreduce:ListInstanceFleets
                  - elasticmapreduce:ListSteps
                  - elasticmapreduce:ListStudios
                  - elasticmapreduce:ListNotebookExecutions
                  - emr-serverless:ListApplications
                  - emr-containers:ListVirtualClusters
                Resource: '*'

  ReviewFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: !Ref ProjectName
      Runtime: python3.12
      Handler: index.handler
      Timeout: 900
      MemorySize: 512
      Role: !GetAtt ReviewRole.Arn
      Environment:
        Variables:
          INCLUDE_DETAILS: !Ref IncludeDetails
          MAX_CLUSTERS_PER_REGION: !Ref MaxClustersPerRegion
          MAX_STEPS_PER_CLUSTER: !Ref MaxStepsPerCluster
          EXCLUDED_REGIONS_CSV: !Ref ExcludedRegionsCsv
      Code:
        ZipFile: |
          import json, os
          from datetime import datetime, timezone
          import boto3
          from botocore.exceptions import ClientError, EndpointConnectionError

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

          def boolish(v):
              return str(v).lower() in ("1","true","yes","y")

          def csv(v):
              return [x.strip() for x in (v or "").split(",") if x.strip()]

          def err(service, op, e):
              if isinstance(e, ClientError):
                  return {"service":service,"operation":op,"error_code":e.response.get("Error",{}).get("Code"),"error_message":e.response.get("Error",{}).get("Message")}
              return {"service":service,"operation":op,"error_code":type(e).__name__,"error_message":str(e)}

          def denied(e):
              s = str(e).lower()
              if isinstance(e, ClientError):
                  s = (e.response.get("Error",{}).get("Code","") + " " + e.response.get("Error",{}).get("Message","")).lower()
              return "accessdenied" in s or "not authorized" in s or "explicit deny" in s or "service control policy" in s

          def safe(rr, svc, op, fn, default):
              try:
                  return fn()
              except EndpointConnectionError as e:
                  rr["failed"] += 1; rr["failures"].append(err(svc,op,e)); return default
              except ClientError as e:
                  if denied(e): rr["access_denied"] += 1
                  else: rr["failed"] += 1
                  rr["failures"].append(err(svc,op,e)); return default
              except Exception as e:
                  rr["failed"] += 1; rr["failures"].append(err(svc,op,e)); return default

          def regions(excluded):
              ec2 = boto3.client("ec2", region_name="us-east-1")
              out = []
              for r in ec2.describe_regions(AllRegions=True).get("Regions", []):
                  st = r.get("OptInStatus")
                  name = r.get("RegionName")
                  if st in (None, "opt-in-not-required", "opted-in") and name not in excluded:
                      out.append(name)
              return sorted(out)

          def list_emr_clusters(emr, max_items):
              states = ["STARTING","BOOTSTRAPPING","RUNNING","WAITING","TERMINATING","TERMINATED","TERMINATED_WITH_ERRORS"]
              items, marker = [], None
              while True:
                  kw = {"ClusterStates": states}
                  if marker: kw["Marker"] = marker
                  resp = emr.list_clusters(**kw)
                  items += resp.get("Clusters", [])
                  if len(items) >= max_items: return items[:max_items]
                  marker = resp.get("Marker")
                  if not marker: return items

          def list_marker(client_call, key, **base):
              items, marker = [], None
              while True:
                  kw = dict(base)
                  if marker: kw["Marker"] = marker
                  resp = client_call(**kw)
                  items += resp.get(key, [])
                  marker = resp.get("Marker")
                  if not marker: return items

          def list_token(client_call, key):
              items, token = [], None
              while True:
                  kw = {}
                  if token: kw["nextToken"] = token
                  resp = client_call(**kw)
                  items += resp.get(key, [])
                  token = resp.get("nextToken")
                  if not token: return items

          def step_counts(steps):
              c = {"sampled_steps":len(steps), "running_steps":0, "pending_steps":0, "failed_steps":0, "cancelled_steps":0, "completed_steps":0, "interrupted_steps":0}
              for s in steps:
                  state = (s.get("Status") or {}).get("State")
                  if state == "RUNNING": c["running_steps"] += 1
                  elif state in ("PENDING","CANCEL_PENDING"): c["pending_steps"] += 1
                  elif state == "FAILED": c["failed_steps"] += 1
                  elif state == "CANCELLED": c["cancelled_steps"] += 1
                  elif state == "COMPLETED": c["completed_steps"] += 1
                  elif state == "INTERRUPTED": c["interrupted_steps"] += 1
              return c

          def group_counts(groups):
              c = {"instance_groups":len(groups), "requested_instances":0, "running_instances":0, "master_groups":0, "core_groups":0, "task_groups":0}
              for g in groups:
                  role = g.get("InstanceGroupType")
                  if role == "MASTER": c["master_groups"] += 1
                  elif role == "CORE": c["core_groups"] += 1
                  elif role == "TASK": c["task_groups"] += 1
                  c["requested_instances"] += int(g.get("RequestedInstanceCount") or 0)
                  c["running_instances"] += int(g.get("RunningInstanceCount") or 0)
              return c

          def scan(region, settings):
              rr = {
                  "region": region, "clusters":0, "active_clusters":0, "running_clusters":0, "waiting_clusters":0,
                  "terminated_clusters":0, "terminated_with_errors_clusters":0, "total_requested_instances":0,
                  "total_running_instances":0, "instance_groups":0, "instance_fleets":0, "sampled_steps":0,
                  "running_steps":0, "pending_steps":0, "failed_steps":0, "cancelled_steps":0, "completed_steps":0,
                  "interrupted_steps":0, "studios":0, "notebook_executions":0, "running_notebook_executions":0,
                  "serverless_applications":0, "serverless_started_applications":0, "serverless_created_applications":0,
                  "serverless_stopped_applications":0, "emr_on_eks_virtual_clusters":0, "emr_on_eks_running_virtual_clusters":0,
                  "review_flags":[], "cluster_inventory":[], "serverless_application_inventory":[], "virtual_cluster_inventory":[],
                  "access_denied":0, "failed":0, "failures":[]
              }
              emr = boto3.client("emr", region_name=region)
              clusters = safe(rr,"EMR","ListClusters",lambda: list_emr_clusters(emr, settings["max_clusters"]),[])
              rr["clusters"] = len(clusters)
              for c in clusters:
                  cid = c.get("Id")
                  detail = safe(rr,"EMR","DescribeCluster",lambda: emr.describe_cluster(ClusterId=cid),{})
                  groups = safe(rr,"EMR","ListInstanceGroups",lambda: list_marker(emr.list_instance_groups,"InstanceGroups",ClusterId=cid),[])
                  fleets = safe(rr,"EMR","ListInstanceFleets",lambda: list_marker(emr.list_instance_fleets,"InstanceFleets",ClusterId=cid),[])
                  steps = safe(rr,"EMR","ListSteps",lambda: list_marker(emr.list_steps,"Steps",ClusterId=cid)[:settings["max_steps"]],[])
                  cluster = detail.get("Cluster",{})
                  state = (cluster.get("Status") or {}).get("State") or (c.get("Status") or {}).get("State")
                  if state in ("STARTING","BOOTSTRAPPING","RUNNING","WAITING"): rr["active_clusters"] += 1
                  if state == "RUNNING": rr["running_clusters"] += 1
                  if state == "WAITING": rr["waiting_clusters"] += 1
                  if state == "TERMINATED": rr["terminated_clusters"] += 1
                  if state == "TERMINATED_WITH_ERRORS": rr["terminated_with_errors_clusters"] += 1
                  gc, sc = group_counts(groups), step_counts(steps)
                  rr["instance_groups"] += gc["instance_groups"]; rr["instance_fleets"] += len(fleets)
                  rr["total_requested_instances"] += gc["requested_instances"]; rr["total_running_instances"] += gc["running_instances"]
                  for k,v in sc.items(): rr[k] += v
                  if settings["details"]:
                      rr["cluster_inventory"].append({"id":cid,"name":cluster.get("Name") or c.get("Name"),"state":state,"release_label":cluster.get("ReleaseLabel"),**gc, "instance_fleets":len(fleets), **sc})
              studios = safe(rr,"EMR","ListStudios",lambda: list_marker(emr.list_studios,"Studios"),[])
              rr["studios"] = len(studios)
              notebooks = safe(rr,"EMR","ListNotebookExecutions",lambda: list_marker(emr.list_notebook_executions,"NotebookExecutions"),[])
              rr["notebook_executions"] = len(notebooks)
              rr["running_notebook_executions"] = len([n for n in notebooks if n.get("Status") in ("STARTING","RUNNING")])
              emrs = boto3.client("emr-serverless", region_name=region)
              apps = safe(rr,"EMRServerless","ListApplications",lambda: list_token(emrs.list_applications,"applications"),[])
              rr["serverless_applications"] = len(apps)
              for a in apps:
                  if a.get("state") == "STARTED": rr["serverless_started_applications"] += 1
                  elif a.get("state") == "CREATED": rr["serverless_created_applications"] += 1
                  elif a.get("state") == "STOPPED": rr["serverless_stopped_applications"] += 1
                  if settings["details"]: rr["serverless_application_inventory"].append(a)
              emrc = boto3.client("emr-containers", region_name=region)
              vcs = safe(rr,"EMROnEKS","ListVirtualClusters",lambda: list_token(emrc.list_virtual_clusters,"virtualClusters"),[])
              rr["emr_on_eks_virtual_clusters"] = len(vcs)
              rr["emr_on_eks_running_virtual_clusters"] = len([v for v in vcs if v.get("state") in ("RUNNING","CREATING")])
              if settings["details"]: rr["virtual_cluster_inventory"] = vcs
              if rr["active_clusters"] > 0: rr["review_flags"].append("active_emr_clusters_present")
              if rr["waiting_clusters"] > 0: rr["review_flags"].append("waiting_emr_clusters_present")
              if rr["total_running_instances"] > 0: rr["review_flags"].append("emr_running_instances_present")
              if rr["running_steps"] > 0 or rr["pending_steps"] > 0: rr["review_flags"].append("active_emr_steps_present")
              if rr["failed_steps"] > 0 or rr["interrupted_steps"] > 0: rr["review_flags"].append("failed_or_interrupted_steps_sampled")
              if rr["serverless_started_applications"] > 0: rr["review_flags"].append("started_emr_serverless_applications_present")
              if rr["emr_on_eks_running_virtual_clusters"] > 0: rr["review_flags"].append("running_emr_on_eks_virtual_clusters_present")
              return rr

          def handler(event, context):
              event = event or {}
              excluded = set(csv(event.get("excluded_regions_csv") or os.environ.get("EXCLUDED_REGIONS_CSV","ap-northeast-3")))
              settings = {
                  "details": boolish(event.get("include_details", os.environ.get("INCLUDE_DETAILS","true"))),
                  "max_clusters": int(event.get("max_clusters_per_region") or os.environ.get("MAX_CLUSTERS_PER_REGION","200")),
                  "max_steps": int(event.get("max_steps_per_cluster") or os.environ.get("MAX_STEPS_PER_CLUSTER","50"))
              }
              totals = {k:0 for k in ["regions_scanned","clusters","active_clusters","running_clusters","waiting_clusters","terminated_clusters","terminated_with_errors_clusters","total_requested_instances","total_running_instances","instance_groups","instance_fleets","sampled_steps","running_steps","pending_steps","failed_steps","cancelled_steps","completed_steps","interrupted_steps","studios","notebook_executions","running_notebook_executions","serverless_applications","serverless_started_applications","serverless_created_applications","serverless_stopped_applications","emr_on_eks_virtual_clusters","emr_on_eks_running_virtual_clusters","regions_with_clusters","regions_with_active_clusters","regions_with_serverless_applications","regions_with_emr_on_eks_virtual_clusters","regions_with_review_flags","access_denied","failed"]}
              out = {"started_at":now(),"review_only":True,"include_details":settings["details"],"excluded_regions":sorted(excluded),"totals":totals,"region_results":[],"failures":[]}
              for r in regions(excluded):
                  rr = scan(r, settings)
                  for k in totals:
                      if k.startswith("regions_with_") or k == "regions_scanned": continue
                      totals[k] += int(rr.get(k,0) or 0)
                  totals["regions_scanned"] += 1
                  if rr.get("clusters",0) > 0: totals["regions_with_clusters"] += 1
                  if rr.get("active_clusters",0) > 0: totals["regions_with_active_clusters"] += 1
                  if rr.get("serverless_applications",0) > 0: totals["regions_with_serverless_applications"] += 1
                  if rr.get("emr_on_eks_virtual_clusters",0) > 0: totals["regions_with_emr_on_eks_virtual_clusters"] += 1
                  if rr.get("review_flags"): totals["regions_with_review_flags"] += 1
                  out["failures"].extend([dict(f, region=r) for f in rr.get("failures",[])])
                  out["region_results"].append(rr)
              out["ended_at"] = now()
              print(json.dumps(out, indent=2, default=str))
              return json.loads(json.dumps(out, default=str))

  ReviewScheduleRule:
    Type: AWS::Events::Rule
    Condition: IsScheduleEnabled
    Properties:
      Name: !Sub '${ProjectName}-schedule'
      ScheduleExpression: !Ref ScheduleExpression
      State: ENABLED
      Targets:
        - Arn: !GetAtt ReviewFunction.Arn
          Id: EmrCostReviewTarget
          Input: '{"source":"eventbridge-schedule"}'

  ReviewSchedulePermission:
    Type: AWS::Lambda::Permission
    Condition: IsScheduleEnabled
    Properties:
      FunctionName: !Ref ReviewFunction
      Action: lambda:InvokeFunction
      Principal: events.amazonaws.com
      SourceArn: !GetAtt ReviewScheduleRule.Arn

Outputs:
  ReviewFunctionName:
    Description: Lambda function name.
    Value: !Ref ReviewFunction
  RemediatorFunctionName:
    Description: Lambda function name alias.
    Value: !Ref ReviewFunction
  ReviewFunctionArn:
    Description: Lambda function ARN.
    Value: !GetAtt ReviewFunction.Arn
  ReviewOnly:
    Description: This module is intentionally review-only.
    Value: 'true'
