# ==============================================================================
# Additional tests for lap_solve_batch to increase coverage
# ==============================================================================

# ------------------------------------------------------------------------------
# Basic input validation
# ------------------------------------------------------------------------------

test_that("lap_solve_batch handles empty list input", {
  expect_error(
    lap_solve_batch(list()),
    "at least one problem"
  )
})

test_that("lap_solve_batch rejects ungrouped data frame with source column", {
  df <- data.frame(source = 1:3, target = 1:3, cost = 1:3)
  expect_error(
    lap_solve_batch(df, source = source, target = target, cost = cost),
    "group_by"
  )
})

test_that("lap_solve_batch rejects invalid input types", {
  expect_error(
    lap_solve_batch("not a valid input"),
    "must be a list"
  )

  expect_error(
    lap_solve_batch(1:10),
    "must be a list"
  )
})

# ------------------------------------------------------------------------------
# 3D array input
# ------------------------------------------------------------------------------

test_that("lap_solve_batch handles 3D array input", {
  arr <- array(runif(2 * 2 * 3), dim = c(2, 2, 3))

  result <- lap_solve_batch(arr)

  expect_s3_class(result, "lap_solve_batch_result")
  expect_true("problem_id" %in% names(result))
  expect_equal(length(unique(result$problem_id)), 3)
})

test_that("lap_solve_batch handles single-slice 3D array", {
  arr <- array(runif(3 * 3), dim = c(3, 3, 1))

  result <- lap_solve_batch(arr)

  expect_equal(length(unique(result$problem_id)), 1)
})

# ------------------------------------------------------------------------------
# List of matrices input
# ------------------------------------------------------------------------------

test_that("lap_solve_batch handles list of matrices", {
  costs <- list(
    matrix(c(1, 2, 3, 4), 2, 2),
    matrix(c(5, 6, 7, 8, 9, 10), 2, 3),
    matrix(c(1, 5, 9, 2, 6, 10, 3, 7, 11), 3, 3)
  )

  result <- lap_solve_batch(costs)

  expect_s3_class(result, "lap_solve_batch_result")
  expect_equal(length(unique(result$problem_id)), 3)
})

test_that("lap_solve_batch handles single matrix in list", {
  costs <- list(matrix(c(1, 2, 3, 4), 2, 2))

  result <- lap_solve_batch(costs)

  expect_equal(length(unique(result$problem_id)), 1)
})

test_that("lap_solve_batch handles matrices with NA values", {
  costs <- list(
    matrix(c(1, NA, 3, 4), 2, 2),
    matrix(c(NA, 2, 3, NA), 2, 2)
  )

  result <- lap_solve_batch(costs)

  expect_s3_class(result, "lap_solve_batch_result")
})

# ------------------------------------------------------------------------------
# Method selection
# ------------------------------------------------------------------------------

test_that("lap_solve_batch respects method parameter", {
  costs <- list(
    matrix(runif(9), 3, 3),
    matrix(runif(9), 3, 3)
  )

  result_hungarian <- lap_solve_batch(costs, method = "hungarian")
  result_jv <- lap_solve_batch(costs, method = "jv")

  expect_true(all(result_hungarian$method_used == "hungarian"))
  expect_true(all(result_jv$method_used == "jv"))
})

test_that("lap_solve_batch handles maximize parameter", {
  costs <- list(
    matrix(c(1, 10, 10, 1), 2, 2),
    matrix(c(1, 10, 10, 1), 2, 2)
  )

  result_min <- lap_solve_batch(costs, maximize = FALSE)
  result_max <- lap_solve_batch(costs, maximize = TRUE)

  expect_true(all(result_min$total_cost < result_max$total_cost))
})

# ------------------------------------------------------------------------------
# Parallel execution
# ------------------------------------------------------------------------------

test_that("lap_solve_batch works with n_threads = 1", {
  costs <- list(
    matrix(runif(9), 3, 3),
    matrix(runif(9), 3, 3)
  )

  result <- lap_solve_batch(costs, n_threads = 1)

  expect_s3_class(result, "lap_solve_batch_result")
})

test_that("lap_solve_batch falls back to sequential for few problems", {
  costs <- list(
    matrix(runif(9), 3, 3),
    matrix(runif(9), 3, 3)
  )

  # With only 2 problems and n_threads = 4, should use sequential
  result <- lap_solve_batch(costs, n_threads = 4)

  expect_s3_class(result, "lap_solve_batch_result")
})

test_that("lap_solve_batch handles n_threads = NULL", {
  skip_on_cran()
  skip_if(nzchar(Sys.getenv("_R_CHECK_LIMIT_CORES_")),
          "parallel tests limited in check environments")

  costs <- list(
    matrix(runif(9), 3, 3),
    matrix(runif(9), 3, 3),
    matrix(runif(9), 3, 3),
    matrix(runif(9), 3, 3),
    matrix(runif(9), 3, 3)
  )

  # n_threads = NULL should detect cores
  result <- lap_solve_batch(costs, n_threads = NULL)

  expect_s3_class(result, "lap_solve_batch_result")
})

# ------------------------------------------------------------------------------
# Grouped data frame input
# ------------------------------------------------------------------------------

test_that("lap_solve_batch handles grouped data frame", {
  df <- tibble::tibble(
    sim = rep(1:3, each = 4),
    source = rep(1:2, times = 6),
    target = rep(1:2, each = 2, times = 3),
    cost = runif(12, 1, 10)
  )

  result <- df |>
    dplyr::group_by(sim) |>
    lap_solve_batch(source, target, cost)

  expect_s3_class(result, "lap_solve_batch_result")
  expect_true("sim" %in% names(result))
  expect_equal(length(unique(result$sim)), 3)
})

test_that("lap_solve_batch grouped requires column specification", {
  df <- tibble::tibble(
    sim = rep(1:2, each = 4),
    source = rep(1:2, times = 4),
    target = rep(1:2, each = 2, times = 2),
    cost = runif(8)
  )

  expect_error(
    df |> dplyr::group_by(sim) |> lap_solve_batch(),
    "must specify"
  )
})

test_that("lap_solve_batch grouped handles forbidden parameter", {
  df <- tibble::tibble(
    sim = rep(1:2, each = 4),
    source = rep(1:2, times = 4),
    target = rep(1:2, each = 2, times = 2),
    cost = c(1, 2, 3, 4, 1, 2, 3, 4)
  )

  result <- df |>
    dplyr::group_by(sim) |>
    lap_solve_batch(source, target, cost, forbidden = Inf)

  expect_s3_class(result, "lap_solve_batch_result")
})

# ------------------------------------------------------------------------------
# Print method
# ------------------------------------------------------------------------------

test_that("print.lap_solve_batch_result works", {
  costs <- list(
    matrix(c(1, 2, 3, 4), 2, 2),
    matrix(c(5, 6, 7, 8), 2, 2)
  )

  result <- lap_solve_batch(costs)

  output <- capture.output(print(result))

  expect_true(any(grepl("Batch Assignment", output)))
  expect_true(any(grepl("Number of problems", output)))
})

test_that("print.lap_solve_batch_result handles missing columns gracefully", {
  # Create a result with missing columns
  result <- tibble::tibble(
    source = 1:2,
    target = c(2, 1),
    cost = c(1.0, 2.0)
  )
  class(result) <- c("lap_solve_batch_result", class(result))

  output <- capture.output(print(result))

  expect_true(any(grepl("Batch Assignment", output)))
})

# ------------------------------------------------------------------------------
# Edge cases
# ------------------------------------------------------------------------------

test_that("lap_solve_batch handles rectangular matrices", {
  costs <- list(
    matrix(runif(6), 2, 3),
    matrix(runif(6), 2, 3)
  )

  result <- lap_solve_batch(costs)

  expect_s3_class(result, "lap_solve_batch_result")
})

test_that("lap_solve_batch handles 1x1 matrices", {
  costs <- list(
    matrix(5, 1, 1),
    matrix(10, 1, 1)
  )

  result <- lap_solve_batch(costs)

  expect_equal(sum(result$total_cost), 15)
})

test_that("lap_solve_batch handles sparse matrices", {
  set.seed(123)
  # Create sparse matrix with guaranteed feasible assignment (diagonal)
  make_sparse_feasible <- function(n) {
    m <- matrix(runif(n * n), n, n)
    # Make ~50% NA but keep diagonal feasible
    for (i in 1:n) {
      for (j in 1:n) {
        if (i != j && runif(1) > 0.5) m[i, j] <- NA
      }
    }
    m
  }

  costs <- list(make_sparse_feasible(5), make_sparse_feasible(5))

  result <- lap_solve_batch(costs)

  expect_s3_class(result, "lap_solve_batch_result")
})

# ------------------------------------------------------------------------------
# Parallel grouped execution
# ------------------------------------------------------------------------------

test_that("lap_solve_batch grouped parallel handles many groups", {
  skip_on_cran()

  df <- tibble::tibble(
    sim = rep(1:10, each = 4),
    source = rep(1:2, times = 20),
    target = rep(1:2, each = 2, times = 10),
    cost = runif(40, 1, 10)
  )

  result <- df |>
    dplyr::group_by(sim) |>
    lap_solve_batch(source, target, cost, n_threads = 2)

  expect_s3_class(result, "lap_solve_batch_result")
  expect_equal(length(unique(result$sim)), 10)
})

# ------------------------------------------------------------------------------
# Additional coverage tests
# ------------------------------------------------------------------------------

test_that("lap_solve_batch parallel execution for many matrices", {
  skip_on_cran()

  # Need >= 4 problems to trigger parallel execution
  costs <- lapply(1:6, function(i) matrix(runif(9), 3, 3))

  result <- lap_solve_batch(costs, n_threads = 2)

  expect_s3_class(result, "lap_solve_batch_result")
  expect_equal(length(unique(result$problem_id)), 6)
})

test_that("lap_solve_batch grouped with n_threads = NULL", {
  skip_on_cran()
  skip_if(nzchar(Sys.getenv("_R_CHECK_LIMIT_CORES_")),
          "parallel tests limited in check environments")

  df <- tibble::tibble(
    sim = rep(1:5, each = 4),
    source = rep(1:2, times = 10),
    target = rep(1:2, each = 2, times = 5),
    cost = runif(20, 1, 10)
  )

  result <- df |>
    dplyr::group_by(sim) |>
    lap_solve_batch(source, target, cost, n_threads = NULL)

  expect_s3_class(result, "lap_solve_batch_result")
})

test_that("lap_solve_batch handles matrices returning empty matches", {
  # Matrix where all entries on diagonal are forbidden (infeasible)
  # but some other assignment is possible
  cost1 <- matrix(c(NA, 1, 1, NA), 2, 2)  # No diagonal assignment, but (1,2),(2,1) works
  cost2 <- matrix(c(1, 2, 3, 4), 2, 2)

  result <- lap_solve_batch(list(cost1, cost2))

  expect_s3_class(result, "lap_solve_batch_result")
})