AWSTemplateFormatVersion: '2010-09-09'
Description: >
  AWS Batch Compute Environment / Job Queue Cost Review Remediator.
  Review-only inventory of AWS Batch compute environments, job queues, and job
  status counts across enabled Regions.

Parameters:
  ProjectName:
    Type: String
    Default: batch-cost-review
  IncludeDetails:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
  MaxJobsPerStatus:
    Type: Number
    Default: 100
    MinValue: 1
    MaxValue: 1000
  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: BatchCostReviewReadOnlyPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action: ec2:DescribeRegions
                Resource: '*'
              - Effect: Allow
                Action:
                  - batch:DescribeComputeEnvironments
                  - batch:DescribeJobQueues
                  - batch:DescribeJobs
                  - batch:ListJobs
                  - batch:DescribeJobDefinitions
                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_JOBS_PER_STATUS: !Ref MaxJobsPerStatus
          EXCLUDED_REGIONS_CSV: !Ref ExcludedRegionsCsv
      Code:
        ZipFile: |
          import json, os
          from datetime import datetime, timezone
          import boto3
          from botocore.exceptions import ClientError, EndpointConnectionError

          JOB_STATUSES = ["SUBMITTED", "PENDING", "RUNNABLE", "STARTING", "RUNNING", "SUCCEEDED", "FAILED"]

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

          def boolish(v):
              return str(v).strip().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 enabled_regions(excluded):
              ec2 = boto3.client("ec2", region_name="us-east-1")
              regions = []
              for r in ec2.describe_regions(AllRegions=True).get("Regions", []):
                  status = r.get("OptInStatus")
                  name = r.get("RegionName")
                  if status in (None, "opt-in-not-required", "opted-in") and name not in excluded:
                      regions.append(name)
              return sorted(regions)

          def describe_all_compute_environments(client):
              out, token = [], None
              while True:
                  kw = {}
                  if token:
                      kw["nextToken"] = token
                  resp = client.describe_compute_environments(**kw)
                  out.extend(resp.get("computeEnvironments", []) or [])
                  token = resp.get("nextToken")
                  if not token:
                      return out

          def describe_all_job_queues(client):
              out, token = [], None
              while True:
                  kw = {}
                  if token:
                      kw["nextToken"] = token
                  resp = client.describe_job_queues(**kw)
                  out.extend(resp.get("jobQueues", []) or [])
                  token = resp.get("nextToken")
                  if not token:
                      return out

          def list_jobs_for_queue_status(client, queue, status, max_items):
              out, token = [], None
              while True:
                  kw = {"jobQueue": queue, "jobStatus": status, "maxResults": min(100, max_items)}
                  if token:
                      kw["nextToken"] = token
                  resp = client.list_jobs(**kw)
                  out.extend(resp.get("jobSummaryList", []) or [])
                  if len(out) >= max_items:
                      return out[:max_items]
                  token = resp.get("nextToken")
                  if not token:
                      return out

          def describe_job_defs(client):
              out, token = [], None
              while True:
                  kw = {"status": "ACTIVE"}
                  if token:
                      kw["nextToken"] = token
                  resp = client.describe_job_definitions(**kw)
                  out.extend(resp.get("jobDefinitions", []) or [])
                  token = resp.get("nextToken")
                  if not token:
                      return out

          def compute_env_type_counts(rr, ce):
              ce_type = ce.get("type")
              state = ce.get("state")
              status = ce.get("status")
              cr = ce.get("computeResources") or {}
              cr_type = cr.get("type")

              if state == "ENABLED":
                  rr["batch_enabled_compute_environments"] += 1
              elif state == "DISABLED":
                  rr["batch_disabled_compute_environments"] += 1
              if status == "VALID":
                  rr["batch_valid_compute_environments"] += 1
              elif status == "INVALID":
                  rr["batch_invalid_compute_environments"] += 1

              if ce_type == "MANAGED":
                  rr["batch_managed_compute_environments"] += 1
              elif ce_type == "UNMANAGED":
                  rr["batch_unmanaged_compute_environments"] += 1

              if cr_type == "EC2":
                  rr["batch_ec2_compute_environments"] += 1
              elif cr_type == "SPOT":
                  rr["batch_spot_compute_environments"] += 1
              elif cr_type == "FARGATE":
                  rr["batch_fargate_compute_environments"] += 1
              elif cr_type == "FARGATE_SPOT":
                  rr["batch_fargate_spot_compute_environments"] += 1

              rr["batch_min_vcpus"] += int(cr.get("minvCpus") or 0)
              rr["batch_desired_vcpus"] += int(cr.get("desiredvCpus") or 0)
              rr["batch_max_vcpus"] += int(cr.get("maxvCpus") or 0)

          def scan(region, details, max_jobs):
              rr = {
                  "region": region,
                  "batch_compute_environments": 0,
                  "batch_enabled_compute_environments": 0,
                  "batch_disabled_compute_environments": 0,
                  "batch_valid_compute_environments": 0,
                  "batch_invalid_compute_environments": 0,
                  "batch_managed_compute_environments": 0,
                  "batch_unmanaged_compute_environments": 0,
                  "batch_ec2_compute_environments": 0,
                  "batch_spot_compute_environments": 0,
                  "batch_fargate_compute_environments": 0,
                  "batch_fargate_spot_compute_environments": 0,
                  "batch_min_vcpus": 0,
                  "batch_desired_vcpus": 0,
                  "batch_max_vcpus": 0,
                  "batch_job_queues": 0,
                  "batch_enabled_job_queues": 0,
                  "batch_disabled_job_queues": 0,
                  "batch_valid_job_queues": 0,
                  "batch_invalid_job_queues": 0,
                  "batch_active_job_definitions": 0,
                  "batch_jobs_submitted_sampled": 0,
                  "batch_jobs_pending_sampled": 0,
                  "batch_jobs_runnable_sampled": 0,
                  "batch_jobs_starting_sampled": 0,
                  "batch_jobs_running_sampled": 0,
                  "batch_jobs_succeeded_sampled": 0,
                  "batch_jobs_failed_sampled": 0,
                  "review_flags": [],
                  "compute_environment_inventory": [],
                  "job_queue_inventory": [],
                  "job_definition_inventory": [],
                  "job_status_samples": [],
                  "access_denied": 0,
                  "failed": 0,
                  "failures": []
              }

              batch = boto3.client("batch", region_name=region)

              ces = safe(rr, "Batch", "DescribeComputeEnvironments", lambda: describe_all_compute_environments(batch), [])
              rr["batch_compute_environments"] = len(ces)
              for ce in ces:
                  compute_env_type_counts(rr, ce)
                  cr = ce.get("computeResources") or {}
                  if details:
                      rr["compute_environment_inventory"].append({
                          "compute_environment_name": ce.get("computeEnvironmentName"),
                          "compute_environment_arn": ce.get("computeEnvironmentArn"),
                          "state": ce.get("state"),
                          "status": ce.get("status"),
                          "status_reason": ce.get("statusReason"),
                          "type": ce.get("type"),
                          "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 [],
                          "subnet_count": len(cr.get("subnets") or []),
                          "security_group_count": len(cr.get("securityGroupIds") or []),
                          "tags": ce.get("tags") or {}
                      })

              queues = safe(rr, "Batch", "DescribeJobQueues", lambda: describe_all_job_queues(batch), [])
              rr["batch_job_queues"] = len(queues)
              for q in queues:
                  if q.get("state") == "ENABLED":
                      rr["batch_enabled_job_queues"] += 1
                  elif q.get("state") == "DISABLED":
                      rr["batch_disabled_job_queues"] += 1
                  if q.get("status") == "VALID":
                      rr["batch_valid_job_queues"] += 1
                  elif q.get("status") == "INVALID":
                      rr["batch_invalid_job_queues"] += 1

                  queue_name = q.get("jobQueueName")
                  if details:
                      rr["job_queue_inventory"].append({
                          "job_queue_name": queue_name,
                          "job_queue_arn": q.get("jobQueueArn"),
                          "state": q.get("state"),
                          "status": q.get("status"),
                          "status_reason": q.get("statusReason"),
                          "priority": q.get("priority"),
                          "compute_environment_order_count": len(q.get("computeEnvironmentOrder") or []),
                          "tags": q.get("tags") or {}
                      })

                  for status in JOB_STATUSES:
                      jobs = safe(rr, "Batch", f"ListJobs:{status}", lambda qn=queue_name, st=status: list_jobs_for_queue_status(batch, qn, st, max_jobs), [])
                      key = "batch_jobs_" + status.lower() + "_sampled"
                      rr[key] += len(jobs)
                      if details and jobs:
                          rr["job_status_samples"].append({
                              "job_queue_name": queue_name,
                              "status": status,
                              "sampled_count": len(jobs),
                              "jobs": [
                                  {
                                      "job_id": j.get("jobId"),
                                      "job_name": j.get("jobName"),
                                      "status": j.get("status"),
                                      "created_at": j.get("createdAt"),
                                      "started_at": j.get("startedAt"),
                                      "stopped_at": j.get("stoppedAt")
                                  }
                                  for j in jobs[:25]
                              ]
                          })

              job_defs = safe(rr, "Batch", "DescribeJobDefinitions", lambda: describe_job_defs(batch), [])
              rr["batch_active_job_definitions"] = len(job_defs)
              if details:
                  rr["job_definition_inventory"] = [
                      {
                          "job_definition_name": jd.get("jobDefinitionName"),
                          "job_definition_arn": jd.get("jobDefinitionArn"),
                          "type": jd.get("type"),
                          "platform_capabilities": jd.get("platformCapabilities") or [],
                          "container_vcpus": (jd.get("containerProperties") or {}).get("vcpus"),
                          "container_memory": (jd.get("containerProperties") or {}).get("memory"),
                          "tags": jd.get("tags") or {}
                      }
                      for jd in job_defs[:200]
                  ]

              if rr["batch_enabled_compute_environments"] > 0:
                  rr["review_flags"].append("enabled_batch_compute_environments_present")
              if rr["batch_desired_vcpus"] > 0:
                  rr["review_flags"].append("batch_desired_vcpu_capacity_present")
              if rr["batch_jobs_running_sampled"] > 0 or rr["batch_jobs_runnable_sampled"] > 0 or rr["batch_jobs_pending_sampled"] > 0:
                  rr["review_flags"].append("active_or_queued_batch_jobs_sampled")
              if rr["batch_invalid_compute_environments"] > 0 or rr["batch_invalid_job_queues"] > 0:
                  rr["review_flags"].append("invalid_batch_resources_present")
              if rr["batch_active_job_definitions"] > 0:
                  rr["review_flags"].append("active_batch_job_definitions_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")))
              details = boolish(event.get("include_details", os.environ.get("INCLUDE_DETAILS", "true")))
              max_jobs = int(event.get("max_jobs_per_status") or os.environ.get("MAX_JOBS_PER_STATUS", "100"))

              total_keys = [
                  "regions_scanned", "batch_compute_environments", "batch_enabled_compute_environments",
                  "batch_disabled_compute_environments", "batch_valid_compute_environments",
                  "batch_invalid_compute_environments", "batch_managed_compute_environments",
                  "batch_unmanaged_compute_environments", "batch_ec2_compute_environments",
                  "batch_spot_compute_environments", "batch_fargate_compute_environments",
                  "batch_fargate_spot_compute_environments", "batch_min_vcpus", "batch_desired_vcpus",
                  "batch_max_vcpus", "batch_job_queues", "batch_enabled_job_queues",
                  "batch_disabled_job_queues", "batch_valid_job_queues", "batch_invalid_job_queues",
                  "batch_active_job_definitions", "batch_jobs_submitted_sampled", "batch_jobs_pending_sampled",
                  "batch_jobs_runnable_sampled", "batch_jobs_starting_sampled", "batch_jobs_running_sampled",
                  "batch_jobs_succeeded_sampled", "batch_jobs_failed_sampled", "regions_with_compute_environments",
                  "regions_with_job_queues", "regions_with_active_or_queued_jobs", "regions_with_review_flags",
                  "access_denied", "failed"
              ]
              totals = {k: 0 for k in total_keys}

              output = {
                  "started_at": now(),
                  "review_only": True,
                  "include_details": details,
                  "excluded_regions": sorted(excluded),
                  "totals": totals,
                  "region_results": [],
                  "failures": []
              }

              for region in enabled_regions(excluded):
                  rr = scan(region, details, max_jobs)
                  for k in total_keys:
                      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["batch_compute_environments"] > 0:
                      totals["regions_with_compute_environments"] += 1
                  if rr["batch_job_queues"] > 0:
                      totals["regions_with_job_queues"] += 1
                  if rr["batch_jobs_running_sampled"] > 0 or rr["batch_jobs_runnable_sampled"] > 0 or rr["batch_jobs_pending_sampled"] > 0:
                      totals["regions_with_active_or_queued_jobs"] += 1
                  if rr["review_flags"]:
                      totals["regions_with_review_flags"] += 1
                  output["failures"].extend([dict(f, region=region) for f in rr.get("failures", [])])
                  output["region_results"].append(rr)

              output["ended_at"] = now()
              print(json.dumps(output, indent=2, default=str))
              return json.loads(json.dumps(output, 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: BatchCostReviewTarget
          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'
