AWSTemplateFormatVersion: '2010-09-09'
Description: AWS Batch Compute 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
  IncludeDetails:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
  IncludeTags:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
  MaxComputeEnvironmentsPerRegion:
    Type: Number
    Default: 100
    MinValue: 1
    MaxValue: 500
  MaxJobQueuesPerRegion:
    Type: Number
    Default: 100
    MinValue: 1
    MaxValue: 500
  MaxJobsPerQueueState:
    Type: Number
    Default: 100
    MinValue: 1
    MaxValue: 500

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

Resources:
  BatchCostReviewRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: [lambda.amazonaws.com]
            Action: [sts:AssumeRole]
      Policies:
        - PolicyName: batch-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: BatchReadOnly
                Effect: Allow
                Action:
                  - batch:DescribeComputeEnvironments
                  - batch:DescribeJobQueues
                  - batch:ListJobs
                  - batch:DescribeJobs
                  - batch:ListTagsForResource
                Resource: '*'

  BatchCostReviewFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: !Sub batch-cost-review-${AWS::StackName}
      Runtime: python3.12
      Handler: index.lambda_handler
      Role: !GetAtt BatchCostReviewRole.Arn
      Timeout: 900
      MemorySize: 256
      Environment:
        Variables:
          EXCLUDED_REGIONS_CSV: !Ref ExcludedRegionsCsv
          INCLUDE_DETAILS: !Ref IncludeDetails
          INCLUDE_TAGS: !Ref IncludeTags
          MAX_COMPUTE_ENVIRONMENTS_PER_REGION: !Ref MaxComputeEnvironmentsPerRegion
          MAX_JOB_QUEUES_PER_REGION: !Ref MaxJobQueuesPerRegion
          MAX_JOBS_PER_QUEUE_STATE: !Ref MaxJobsPerQueueState
      Code:
        ZipFile: |
          import boto3
          import os
          import traceback
          from datetime import datetime, timezone
          from botocore.exceptions import ClientError, EndpointConnectionError

          JOB_STATES = ["SUBMITTED", "PENDING", "RUNNABLE", "STARTING", "RUNNING"]

          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)
              regions = []
              for r in response.get("Regions", []):
                  status = r.get("OptInStatus")
                  if status in (None, "opt-in-not-required", "opted-in"):
                      regions.append(r["RegionName"])
              return sorted(regions)

          def list_compute_environments(batch, max_items):
              items = []
              token = None
              while True:
                  kwargs = {"maxResults": min(max_items, 100)}
                  if token:
                      kwargs["nextToken"] = token
                  response = batch.describe_compute_environments(**kwargs)
                  items.extend(response.get("computeEnvironments", []) or [])
                  if len(items) >= max_items:
                      return items[:max_items], True
                  token = response.get("nextToken")
                  if not token:
                      return items, False

          def list_job_queues(batch, max_items):
              items = []
              token = None
              while True:
                  kwargs = {"maxResults": min(max_items, 100)}
                  if token:
                      kwargs["nextToken"] = token
                  response = batch.describe_job_queues(**kwargs)
                  items.extend(response.get("jobQueues", []) or [])
                  if len(items) >= max_items:
                      return items[:max_items], True
                  token = response.get("nextToken")
                  if not token:
                      return items, False

          def list_jobs_for_queue(batch, queue_name, max_items):
              counts = {state.lower(): 0 for state in JOB_STATES}
              sampled = []
              truncated = False

              for state in JOB_STATES:
                  token = None
                  state_count = 0
                  while True:
                      kwargs = {
                          "jobQueue": queue_name,
                          "jobStatus": state,
                          "maxResults": min(max_items, 100)
                      }
                      if token:
                          kwargs["nextToken"] = token
                      response = batch.list_jobs(**kwargs)
                      summaries = response.get("jobSummaryList", []) or []
                      state_count += len(summaries)
                      for s in summaries[:10]:
                          sampled.append({
                              "job_id": s.get("jobId"),
                              "job_name": s.get("jobName"),
                              "status": s.get("status"),
                              "created_at": s.get("createdAt"),
                              "started_at": s.get("startedAt")
                          })
                      if state_count >= max_items:
                          truncated = True
                          break
                      token = response.get("nextToken")
                      if not token:
                          break
                  counts[state.lower()] = state_count

              return counts, sampled[:50], truncated

          def get_tags(batch, arn):
              if not arn:
                  return {}, None
              try:
                  response = batch.list_tags_for_resource(resourceArn=arn)
                  return response.get("tags", {}) or {}, None
              except ClientError as e:
                  return {}, e.response.get("Error", {}).get("Code", "ClientError")
              except Exception as e:
                  return {}, str(e)

          def summarize_compute_environment(batch, ce, include_tags):
              cr = ce.get("computeResources") or {}
              arn = ce.get("computeEnvironmentArn")
              item = {
                  "compute_environment_name": ce.get("computeEnvironmentName"),
                  "compute_environment_arn": arn,
                  "type": ce.get("type"),
                  "state": ce.get("state"),
                  "status": ce.get("status"),
                  "status_reason": ce.get("statusReason"),
                  "service_role_present": bool(ce.get("serviceRole")),
                  "unmanagedv_cpus": ce.get("unmanagedvCpus"),
                  "ecs_cluster_arn_present": bool(ce.get("ecsClusterArn")),
                  "compute_resource_type": cr.get("type"),
                  "allocation_strategy": cr.get("allocationStrategy"),
                  "min_vcpus": cr.get("minvCpus"),
                  "desired_vcpus": cr.get("desiredvCpus"),
                  "max_vcpus": cr.get("maxvCpus"),
                  "instance_types": cr.get("instanceTypes") or [],
                  "subnets_count": len(cr.get("subnets") or []),
                  "security_group_count": len(cr.get("securityGroupIds") or []),
                  "ec2_key_pair_present": bool(cr.get("ec2KeyPair")),
                  "instance_role_present": bool(cr.get("instanceRole")),
                  "bid_percentage": cr.get("bidPercentage"),
                  "spot_iam_fleet_role_present": bool(cr.get("spotIamFleetRole")),
                  "launch_template_present": bool(cr.get("launchTemplate")),
                  "tags_count_in_compute_resources": len(cr.get("tags") or {}),
                  "review_flags": []
              }

              if include_tags:
                  tags, error = get_tags(batch, arn)
                  item["tags"] = tags
                  item["tag_count"] = len(tags)
                  if error:
                      item["tags_error"] = error

              flags = []
              if item.get("state") == "ENABLED":
                  flags.append("enabled_compute_environment_review")
              if item.get("status") and item.get("status") != "VALID":
                  flags.append("compute_environment_status_" + str(item.get("status")).lower())
              if item.get("type") == "MANAGED":
                  flags.append("managed_compute_environment_review")
              if item.get("compute_resource_type") in ("EC2", "SPOT"):
                  flags.append("ec2_or_spot_compute_resources_review")
              if item.get("compute_resource_type") in ("FARGATE", "FARGATE_SPOT"):
                  flags.append("fargate_compute_resources_review")
              if item.get("desired_vcpus", 0) and item.get("desired_vcpus") > 0:
                  flags.append("desired_vcpus_nonzero_review")
              if item.get("max_vcpus", 0) and item.get("max_vcpus") >= 64:
                  flags.append("high_max_vcpus_review")
              if item.get("min_vcpus", 0) and item.get("min_vcpus") > 0:
                  flags.append("min_vcpus_nonzero_review")
              if item.get("bid_percentage"):
                  flags.append("spot_bid_percentage_review")

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

          def summarize_job_queue(batch, jq, include_tags, max_jobs):
              arn = jq.get("jobQueueArn")
              name = jq.get("jobQueueName")
              ce_order = jq.get("computeEnvironmentOrder") or []
              item = {
                  "job_queue_name": name,
                  "job_queue_arn": arn,
                  "state": jq.get("state"),
                  "status": jq.get("status"),
                  "status_reason": jq.get("statusReason"),
                  "priority": jq.get("priority"),
                  "compute_environment_order_count": len(ce_order),
                  "compute_environment_names": [x.get("computeEnvironment") for x in ce_order],
                  "review_flags": []
              }

              try:
                  counts, sampled, truncated = list_jobs_for_queue(batch, name, max_jobs)
                  item["job_counts"] = counts
                  item["sampled_jobs"] = sampled
                  item["jobs_truncated"] = truncated
              except ClientError as e:
                  item["jobs_error"] = e.response.get("Error", {}).get("Code", "ClientError")
                  item["job_counts"] = {}
              except Exception as e:
                  item["jobs_error"] = str(e)
                  item["job_counts"] = {}

              if include_tags:
                  tags, error = get_tags(batch, arn)
                  item["tags"] = tags
                  item["tag_count"] = len(tags)
                  if error:
                      item["tags_error"] = error

              flags = []
              if item.get("state") == "ENABLED":
                  flags.append("enabled_job_queue_review")
              if item.get("status") and item.get("status") != "VALID":
                  flags.append("job_queue_status_" + str(item.get("status")).lower())
              counts = item.get("job_counts") or {}
              if (counts.get("runnable") or 0) > 0:
                  flags.append("runnable_jobs_review")
              if (counts.get("running") or 0) > 0:
                  flags.append("running_jobs_review")
              if (counts.get("submitted") or 0) + (counts.get("pending") or 0) + (counts.get("starting") or 0) > 0:
                  flags.append("queued_jobs_review")
              if item.get("compute_environment_order_count", 0) == 0:
                  flags.append("job_queue_without_compute_environment_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_DETAILS", "true")), True)
              include_tags = as_bool(event.get("include_tags", os.environ.get("INCLUDE_TAGS", "true")), True)
              max_ces = as_int(event.get("max_compute_environments_per_region", os.environ.get("MAX_COMPUTE_ENVIRONMENTS_PER_REGION", "100")), 100)
              max_queues = as_int(event.get("max_job_queues_per_region", os.environ.get("MAX_JOB_QUEUES_PER_REGION", "100")), 100)
              max_jobs = as_int(event.get("max_jobs_per_queue_state", os.environ.get("MAX_JOBS_PER_QUEUE_STATE", "100")), 100)

              result = {
                  "module": "aws_batch_compute_environment_cost_review",
                  "review_only": True,
                  "started_at": now(),
                  "ended_at": None,
                  "excluded_regions": sorted(excluded),
                  "job_states_sampled": JOB_STATES,
                  "totals": {
                      "regions_scanned": 0,
                      "regions_skipped": 0,
                      "compute_environments": 0,
                      "managed_compute_environments": 0,
                      "enabled_compute_environments": 0,
                      "ec2_or_spot_compute_environments": 0,
                      "fargate_compute_environments": 0,
                      "desired_vcpus": 0,
                      "max_vcpus": 0,
                      "job_queues": 0,
                      "enabled_job_queues": 0,
                      "submitted_jobs": 0,
                      "pending_jobs": 0,
                      "runnable_jobs": 0,
                      "starting_jobs": 0,
                      "running_jobs": 0,
                      "regions_with_batch_resources": 0,
                      "regions_with_review_flags": 0,
                      "truncated_regions": 0,
                      "access_denied": 0,
                      "failed": 0
                  },
                  "regions": [],
                  "compute_environments": [],
                  "job_queues": []
              }

              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,
                      "compute_environments": 0,
                      "managed_compute_environments": 0,
                      "enabled_compute_environments": 0,
                      "ec2_or_spot_compute_environments": 0,
                      "fargate_compute_environments": 0,
                      "desired_vcpus": 0,
                      "max_vcpus": 0,
                      "job_queues": 0,
                      "enabled_job_queues": 0,
                      "submitted_jobs": 0,
                      "pending_jobs": 0,
                      "runnable_jobs": 0,
                      "starting_jobs": 0,
                      "running_jobs": 0,
                      "review_flags": [],
                      "truncated": False,
                      "error": None
                  }

                  try:
                      batch = boto3.client("batch", region_name=region)
                      ces, ce_trunc = list_compute_environments(batch, max_ces)
                      queues, q_trunc = list_job_queues(batch, max_queues)
                      rec["truncated"] = ce_trunc or q_trunc
                      if rec["truncated"]:
                          result["totals"]["truncated_regions"] += 1

                      ce_details = []
                      for ce in ces:
                          d = summarize_compute_environment(batch, ce, include_tags)
                          d["region"] = region
                          ce_details.append(d)
                          rec["compute_environments"] += 1
                          result["totals"]["compute_environments"] += 1

                          if d.get("type") == "MANAGED":
                              rec["managed_compute_environments"] += 1
                              result["totals"]["managed_compute_environments"] += 1
                          if d.get("state") == "ENABLED":
                              rec["enabled_compute_environments"] += 1
                              result["totals"]["enabled_compute_environments"] += 1
                          if d.get("compute_resource_type") in ("EC2", "SPOT"):
                              rec["ec2_or_spot_compute_environments"] += 1
                              result["totals"]["ec2_or_spot_compute_environments"] += 1
                          if d.get("compute_resource_type") in ("FARGATE", "FARGATE_SPOT"):
                              rec["fargate_compute_environments"] += 1
                              result["totals"]["fargate_compute_environments"] += 1

                          desired = d.get("desired_vcpus") or 0
                          maxv = d.get("max_vcpus") or 0
                          rec["desired_vcpus"] += desired
                          rec["max_vcpus"] += maxv
                          result["totals"]["desired_vcpus"] += desired
                          result["totals"]["max_vcpus"] += maxv
                          rec["review_flags"].extend(d.get("review_flags", []))

                      q_details = []
                      for jq in queues:
                          d = summarize_job_queue(batch, jq, include_tags, max_jobs)
                          d["region"] = region
                          q_details.append(d)
                          rec["job_queues"] += 1
                          result["totals"]["job_queues"] += 1

                          if d.get("state") == "ENABLED":
                              rec["enabled_job_queues"] += 1
                              result["totals"]["enabled_job_queues"] += 1

                          counts = d.get("job_counts") or {}
                          for key in ["submitted", "pending", "runnable", "starting", "running"]:
                              value = counts.get(key) or 0
                              rec[key + "_jobs"] += value
                              result["totals"][key + "_jobs"] += value

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

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

                      if rec["compute_environments"] > 0 or rec["job_queues"] > 0:
                          result["totals"]["regions_with_batch_resources"] += 1
                      if rec["review_flags"]:
                          result["totals"]["regions_with_review_flags"] += 1

                      if include_details:
                          result["compute_environments"].extend(ce_details)
                          result["job_queues"].extend(q_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

  BatchCostReviewSchedule:
    Type: AWS::Events::Rule
    Condition: CreateSchedule
    Properties:
      Description: Scheduled review-only AWS Batch compute environment cost review.
      ScheduleExpression: !Ref ScheduleExpression
      State: ENABLED
      Targets:
        - Arn: !GetAtt BatchCostReviewFunction.Arn
          Id: BatchCostReviewFunctionTarget
          Input: '{"scheduled":true}'

  BatchCostReviewSchedulePermission:
    Type: AWS::Lambda::Permission
    Condition: CreateSchedule
    Properties:
      FunctionName: !Ref BatchCostReviewFunction
      Action: lambda:InvokeFunction
      Principal: events.amazonaws.com
      SourceArn: !GetAtt BatchCostReviewSchedule.Arn

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