test_that("session state functions work correctly", {
  skip_on_cran()
  skip_if_not(check_aws_credentials(), "AWS credentials not available")

  # Setup
  config <- get_starburst_config()
  session_id <- sprintf("test-session-%s", gsub("-", "", uuid::UUIDgenerate()))
  region <- config$region %||% "us-east-1"
  bucket <- config$bucket

  # Create minimal backend
  backend <- list(
    session_id = session_id,
    workers = 2,
    cpu = 1,
    memory = "2GB",
    region = region,
    bucket = bucket,
    timeout = 600,
    absolute_timeout = 3600,
    image_uri = "dummy-image",
    cluster = "test-cluster",
    cluster_name = "test-cluster",
    launch_type = "FARGATE",
    instance_type = NULL,
    use_spot = FALSE,
    architecture = "X86_64",
    warm_pool_timeout = 3600,
    capacity_provider_name = "test-provider",
    asg_name = "test-asg",
    aws_account_id = "123456789012",
    subnets = list("subnet-1"),
    security_groups = list("sg-1"),
    task_definition_arn = "arn:aws:ecs:us-east-1:123456789012:task-definition/test:1"
  )

  # Test create_session_manifest
  expect_silent(create_session_manifest(session_id, backend))

  # Test get_session_manifest
  manifest <- get_session_manifest(session_id, region, bucket)
  expect_equal(manifest$session_id, session_id)
  expect_equal(manifest$backend$workers, 2)
  expect_equal(manifest$stats$total_tasks, 0)

  # Test update_session_manifest
  expect_silent(update_session_manifest(
    session_id = session_id,
    updates = list(stats = list(total_tasks = 5, pending = 3)),
    region = region,
    bucket = bucket
  ))

  manifest <- get_session_manifest(session_id, region, bucket)
  expect_equal(manifest$stats$total_tasks, 5)
  expect_equal(manifest$stats$pending, 3)

  # Test create_task_status
  task_id <- "test-task-1"
  expect_silent(create_task_status(session_id, task_id, "pending", region, bucket))

  # Test get_task_status
  status <- get_task_status(session_id, task_id, region, bucket)
  expect_equal(status$task_id, task_id)
  expect_equal(status$state, "pending")
  expect_null(status$claimed_by)

  # Test atomic_claim_task
  worker_id <- "worker-1"
  claimed <- atomic_claim_task(session_id, task_id, worker_id, region, bucket)
  expect_true(claimed)

  status <- get_task_status(session_id, task_id, region, bucket)
  expect_equal(status$state, "claimed")
  expect_equal(status$claimed_by, worker_id)

  # Test that second claim fails (atomic)
  claimed2 <- atomic_claim_task(session_id, task_id, "worker-2", region, bucket)
  expect_false(claimed2)

  # Test list_pending_tasks
  task_id2 <- "test-task-2"
  create_task_status(session_id, task_id2, "pending", region, bucket)

  pending <- list_pending_tasks(session_id, region, bucket)
  expect_true(task_id2 %in% pending)
  expect_false(task_id %in% pending)  # Already claimed

  # Test list_task_statuses
  statuses <- list_task_statuses(session_id, region, bucket)
  expect_equal(length(statuses), 2)
  expect_true(task_id %in% names(statuses))
  expect_true(task_id2 %in% names(statuses))

  # Cleanup
  s3 <- get_s3_client(region)
  prefix <- sprintf("sessions/%s/", session_id)
  tryCatch({
    objects <- s3$list_objects_v2(Bucket = bucket, Prefix = prefix)
    if (length(objects$Contents) > 0) {
      for (obj in objects$Contents) {
        s3$delete_object(Bucket = bucket, Key = obj$Key)
      }
    }
  }, error = function(e) {
    # Ignore cleanup errors
  })
})

test_that("detached session creation works", {
  skip_if_no_integration()  # launches real Fargate workers
  skip_on_cran()
  skip_if_not(check_aws_credentials(), "AWS credentials not available")

  # This test launches actual workers
  session <- starburst_session(
    workers = 2,
    cpu = 1,
    memory = "2GB",
    launch_type = "FARGATE"
  )

  expect_s3_class(session, "StarburstSession")
  expect_true(!is.null(session$session_id))
  expect_true(!is.null(session$backend))

  # Cleanup
  session$cleanup()
})

test_that("session submit and collect workflow", {
  skip_if_no_integration()  # launches real Fargate workers
  skip_on_cran()
  skip_if_not(check_aws_credentials(), "AWS credentials not available")

  # Create session (FARGATE - no EC2 ASG required)
  session <- starburst_session(workers = 2, cpu = 1, memory = "2GB", launch_type = "FARGATE")

  # Submit tasks
  task_ids <- lapply(1:5, function(i) {
    session$submit(quote(i * 2))
  })

  expect_equal(length(task_ids), 5)

  # Check status
  status <- session$status()
  expect_s3_class(status, "StarburstSessionStatus")
  expect_equal(status$total, 5)

  # Wait and collect (Fargate cold start takes ~2-3 min)
  results <- session$collect(wait = TRUE, timeout = 300)

  expect_true(length(results) > 0)

  # Cleanup
  session$cleanup()
})

test_that("session attach workflow", {
  skip_if_no_integration()  # launches real Fargate workers
  skip_on_cran()
  skip_if_not(check_aws_credentials(), "AWS credentials not available")

  # Create session (FARGATE - no EC2 ASG required)
  session <- starburst_session(workers = 2, cpu = 1, memory = "2GB", launch_type = "FARGATE")
  session_id <- session$session_id

  # Submit tasks
  lapply(1:10, function(i) session$submit(quote({Sys.sleep(5); i})))

  # Detach (simulate closing R)
  rm(session)

  # Reattach
  session <- starburst_session_attach(session_id)

  expect_s3_class(session, "StarburstSession")
  expect_equal(session$session_id, session_id)

  # Check status
  status <- session$status()
  expect_equal(status$total, 10)

  # Collect results (Fargate cold start takes ~2-3 min)
  results <- session$collect(wait = TRUE, timeout = 300)
  expect_equal(length(results), 10)

  # Cleanup
  session$cleanup()
})

test_that("list sessions works", {
  skip_on_cran()
  skip_if_not(check_aws_credentials(), "AWS credentials not available")

  sessions <- starburst_list_sessions()

  expect_true(is.data.frame(sessions))
  expect_true("session_id" %in% colnames(sessions))
})

test_that("worker detached mode detection works", {
  # This tests the logic without actually running a worker

  # Mock task with session_id (detached mode)
  detached_task <- list(
    session_id = "session-123",
    task_id = "task-456",
    expr = quote(2 + 2),
    globals = list(),
    packages = NULL
  )

  expect_true(!is.null(detached_task$session_id))

  # Mock task without session_id (ephemeral mode)
  ephemeral_task <- list(
    task_id = "task-789",
    expr = quote(3 + 3),
    globals = list(),
    packages = NULL
  )

  expect_true(is.null(ephemeral_task$session_id))
})

test_that("plan guard prevents detached mode misuse", {
  expect_error(
    plan.starburst(
      strategy = starburst,
      workers = 10,
      detached = TRUE
    ),
    "Detached mode cannot be used with plan"
  )
})

test_that("collect(wait=TRUE) returns when user tasks finish despite live bootstrap tasks", {
  # Regression for #38: the all_completed check must ignore bootstrap-* tasks,
  # which stay "claimed"/"running" for the session's life. Otherwise wait=TRUE
  # hangs until timeout even though every real result has landed in S3.
  skip_if_not_installed("mockery")

  statuses <- list(
    "bootstrap-session-x-1" = list(state = "claimed"),
    "bootstrap-session-x-2" = list(state = "running"),
    "task-aaa" = list(state = "completed"),
    "task-bbb" = list(state = "completed")
  )

  mockery::stub(collect_session_results, "list_task_statuses", function(...) statuses)
  mockery::stub(collect_session_results, "get_s3_client", function(...) {
    list(download_file = function(Bucket, Key, Filename) {
      qs2::qs_save(list(error = FALSE, value = 42), Filename)
    })
  })

  session <- list(backend = list(session_id = "session-x", region = "us-east-1",
                                 bucket = "b"))

  # Must return quickly (well under the timeout) with both user results, not hang.
  start <- Sys.time()
  res <- collect_session_results(session, wait = TRUE, timeout = 30)
  elapsed <- as.numeric(difftime(Sys.time(), start, units = "secs"))

  expect_named(res, c("task-aaa", "task-bbb"), ignore.order = TRUE)
  expect_lt(elapsed, 10)  # returned promptly, did not spin to timeout
})

test_that("collect() returns a structured failure entry for failed tasks", {
  # Regression for the 4th-review finding: collect() must NOT silently drop
  # failed tasks. Every terminal task appears; failures carry error=TRUE.
  skip_if_not_installed("mockery")

  statuses <- list(
    "bootstrap-session-x-1" = list(state = "claimed"),
    "task-ok"   = list(state = "completed"),
    "task-fail" = list(state = "failed", error = "boom")
  )
  mockery::stub(collect_session_results, "list_task_statuses", function(...) statuses)
  # S3 has the success blob but NOT the failed one -> exercises the synthesize path.
  mockery::stub(collect_session_results, "get_s3_client", function(...) {
    list(download_file = function(Bucket, Key, Filename) {
      if (grepl("task-ok", Key)) {
        qs2::qs_save(list(error = FALSE, value = 21), Filename)
      } else {
        stop("NoSuchKey")   # failed task's blob missing
      }
    })
  })

  session <- list(backend = list(session_id = "session-x", region = "us-east-1",
                                 bucket = "b"))
  res <- collect_session_results(session, wait = TRUE, timeout = 30)

  expect_setequal(names(res), c("task-ok", "task-fail"))
  expect_false(isTRUE(res[["task-ok"]]$error))
  expect_equal(res[["task-ok"]]$value, 21)
  # failed task present as a structured failure, synthesized from status$error
  expect_true(isTRUE(res[["task-fail"]]$error))
  expect_match(res[["task-fail"]]$message, "boom")
})
