#!/usr/bin/env nu
# Inspect an ECS service: show task status, or fetch recent CloudWatch logs.
# Requires: AWS CLI v2, configured credentials, and permissions for ECS/CloudWatch Logs.
#
# Examples:
#   ecs tasks my-cluster my-service
#   ecs logs  my-cluster my-service --tail 100
#   ecs logs  my-cluster my-service --region us-east-1

const default_tail = 50

def main [
  mode: string # `tasks` or `logs`
  cluster: string # ECS cluster name or ARN
  service: string # ECS service name or ARN
  --tail: int = $default_tail # Number of recent log events to show
  --region: string # AWS region; otherwise AWS CLI configuration is used
] {
  if $mode not-in [tasks logs] {
    error make {msg: $"mode must be `tasks` or `logs`; got `($mode)`"}
  }

  let aws_args = if ($region | is-empty) { [] } else { [--region $region] }
  let task_arns = (
    ^aws ...$aws_args ecs list-tasks
      --cluster $cluster
      --service-name $service
      --desired-status RUNNING
      --query 'taskArns'
      --output json
    | from json
  )

  if ($task_arns | is-empty) {
    print $"No RUNNING tasks found for service `($service)` in cluster `($cluster)`."
    return
  }

  if $mode == tasks {
    (
      ^aws ...$aws_args ecs describe-tasks
        --cluster $cluster
        --tasks ...$task_arns
        --query 'tasks[].{taskArn:taskArn,lastStatus:lastStatus,desiredStatus:desiredStatus,healthStatus:healthStatus,startedAt:startedAt,stoppedReason:stoppedReason,containers:containers[].{name:name,lastStatus:lastStatus,exitCode:exitCode,reason:reason}}'
        --output json
      | from json
      | table -e
      | print
    )
    return
  }

  let definition_arn = (
    ^aws ...$aws_args ecs describe-tasks
      --cluster $cluster
      --tasks ($task_arns | first)
      --query 'tasks[0].taskDefinitionArn'
      --output text
    | str trim
  )

  let container_definitions = (
    ^aws ...$aws_args ecs describe-task-definition
      --task-definition $definition_arn
      --query 'taskDefinition.containerDefinitions[].{name:name,logDriver:logConfiguration.logDriver,options:logConfiguration.options}'
      --output json
    | from json
  )

  let awslogs = ($container_definitions | where logDriver == 'awslogs')
  if ($awslogs | is-empty) {
    error make {msg: $"Task definition `($definition_arn)` has no containers using the awslogs log driver."}
  }

  for container in $awslogs {
    let group = ($container.options | get 'awslogs-group')
    let prefix = ($container.options | get 'awslogs-stream-prefix')
    let stream_prefix = $"($prefix)/($container.name)/"

    # get-log-events returns a stream's most recent events by default, so walk
    # streams newest-first and stop once no remaining stream can hold an event
    # newer than the ones collected. (describe-log-streams can't combine a
    # stream prefix with LastEventTime ordering, so filter by prefix here.)
    print $"\n==> CloudWatch log group: ($group), stream prefix: ($stream_prefix)"
    let limit = ([$tail 10000] | math min)
    mut events = []
    mut next_token = ''
    mut done = false
    loop {
      let token_args = if ($next_token | is-empty) { [] } else { [--next-token $next_token] }
      let page = (
        ^aws ...$aws_args logs describe-log-streams
          --log-group-name $group
          --order-by LastEventTime
          --descending
          --limit 50
          --output json
          ...$token_args
        | from json
      )
      # project the fields we need and keep the page's nextToken
      let streams = (
        $page.logStreams?
        | default []
        | each {|s| {name: $s.logStreamName, last: ($s.lastEventTimestamp? | default null)}}
        | where name starts-with $stream_prefix
      )
      let next = ($page.nextToken? | default '')
      for stream in $streams {
        let last = ($stream.last? | default null)
        let have_enough = ($events | length) >= $tail
        let oldest = (if ($events | is-empty) { null } else { $events | get timestamp | math min })
        # skip empty streams; stop when this stream's newest event is older
        # than what we already have (later streams are even older)
        if $last == null or ($have_enough and $last < $oldest) {
          if $last != null {
            $done = true
            break
          }
          continue
        }
        let recent = (
          ^aws ...$aws_args logs get-log-events
            --log-group-name $group
            --log-stream-name $stream.name
            --limit $limit
            --query 'events[].{timestamp:timestamp,message:message}'
            --output json
          | from json
          | each {|e| {timestamp: $e.timestamp, message: $e.message, stream: $stream.name}}
        )
        $events = ($events | append $recent | sort-by timestamp | last $tail)
      }
      if $done or ($next | is-empty) {
        break
      }
      $next_token = $next
    }

    for event in $events {
      let time = (($event.timestamp | into int) * 1000000 | into datetime | format date '%Y-%m-%dT%H:%M:%S%.3f%:z')
      let task_id = ($event.stream | split row '/' | last)
      print $"($time)  ($task_id)  ($event.message)"
    }
  }
}