# test_main.R
# Unit tests for mtap-main.R functions

library(testthat)
library(MTaP)

# Helper function to create test files (similar to test_config.R)
create_test_files <- function() {
  # Create sample OTU table
  otu_file <- tempfile(fileext = ".csv")
  otu_data <- data.frame(
    OTU = c("OTU_1", "OTU_2", "OTU_3"),
    Sample1 = c(100, 50, 25),
    Sample2 = c(80, 30, 15),
    Sample3 = c(120, 40, 30),
    stringsAsFactors = FALSE
  )
  write.csv(otu_data, otu_file, row.names = FALSE)
  
  # Create sample taxonomy table
  tax_file <- tempfile(fileext = ".csv")
  tax_data <- data.frame(
    OTU = c("OTU_1", "OTU_2", "OTU_3"),
    Kingdom = c("Bacteria", "Bacteria", "Bacteria"),
    Phylum = c("Firmicutes", "Bacteroidetes", "Proteobacteria"),
    Class = c("Clostridia", "Bacteroidia", "Gammaproteobacteria"),
    Order = c("Clostridiales", "Bacteroidales", "Enterobacterales"),
    Family = c("Lachnospiraceae", "Bacteroidaceae", "Enterobacteriaceae"),
    Genus = c("Blautia", "Bacteroides", "Escherichia"),
    Species = c("Blautia_faecis", "Bacteroides_fragilis", "Escherichia_coli"),
    stringsAsFactors = FALSE
  )
  write.csv(tax_data, tax_file, row.names = FALSE)
  
  return(list(otu_file = otu_file, tax_file = tax_file))
}

# Helper function to check if files exist
check_output_files <- function(prefix, levels) {
  files_exist <- sapply(levels, function(level) {
    file_path <- paste0(prefix, "_", tolower(level), ".csv")
    file.exists(file_path)
  })
  summary_file <- paste0(prefix, "_summary.txt")
  return(c(files_exist, summary = file.exists(summary_file)))
}

# Helper function to clean up test files
cleanup_test_files <- function(files) {
  if (!is.null(files$otu_file) && file.exists(files$otu_file)) {
    unlink(files$otu_file)
  }
  if (!is.null(files$tax_file) && file.exists(files$tax_file)) {
    unlink(files$tax_file)
  }
}

# Helper function to clean up output files
cleanup_output_files <- function(prefix, levels) {
  # Remove generated output files
  for (level in levels) {
    file_path <- paste0(prefix, "_", tolower(level), ".csv")
    if (file.exists(file_path)) {
      unlink(file_path)
    }
  }
  summary_file <- paste0(prefix, "_summary.txt")
  if (file.exists(summary_file)) {
    unlink(summary_file)
  }
}

# =============================================
# Tests for mtap_analyze()
# =============================================

test_that("mtap_analyze works with basic configuration", {
  test_files <- create_test_files()
  
  tryCatch({
    # Create configuration
    config <- mtap_config(
      otu_file = test_files$otu_file,
      tax_file = test_files$tax_file,
      output_prefix = "test_output"
    )
    
    # Run analysis (capture output to avoid cluttering test output)
    capture.output({
      result <- suppressMessages(mtap_analyze(config))
    }, type = "output")
    
    # Test 1: Returns correct class
    expect_s3_class(result, "mtap_result")
    expect_type(result, "list")
    
    # Test 2: Contains expected components
    expected_components <- c(
      "config", "otu_data", "tax_data", "merged_data",
      "taxonomy_levels", "abundance_tables", "summary_file",
      "output_files"
    )
    expect_named(result, expected_components, ignore.order = TRUE)
    
    # Test 3: Check basic structure
    expect_s3_class(result$config, "mtap_config")
    expect_type(result$otu_data, "list")
    expect_type(result$tax_data, "list")
    expect_type(result$merged_data, "list")
    expect_type(result$taxonomy_levels, "character")
    expect_type(result$abundance_tables, "list")
    expect_type(result$summary_file, "character")
    
    # Test 4: Check output files structure
    expect_named(result$output_files, c("abundance", "summary"))
    expect_type(result$output_files$abundance, "character")
    expect_type(result$output_files$summary, "character")
    
    # Test 5: Check that files were created
    default_levels <- c("Kingdom", "Phylum", "Class", "Order", "Family", "Genus", "Species")
    file_checks <- check_output_files("test_output", default_levels)
    expect_true(all(file_checks), 
                info = paste("Missing files:", 
                             paste(names(file_checks)[!file_checks], collapse = ", ")))
    
    # Clean up output files
    cleanup_output_files("test_output", default_levels)
    
  }, finally = {
    # Clean up test files
    cleanup_test_files(test_files)
  })
})

test_that("mtap_analyze works with custom taxonomy hierarchy", {
  test_files <- create_test_files()
  
  tryCatch({
    # Create configuration with custom hierarchy
    config <- mtap_config(
      otu_file = test_files$otu_file,
      tax_file = test_files$tax_file,
      output_prefix = "test_custom",
      taxonomy_hierarchy = c("Phylum", "Class", "Genus")
    )
    
    # Run analysis
    capture.output({
      result <- suppressMessages(mtap_analyze(config))
    }, type = "output")
    
    # Test 1: Uses custom taxonomy levels
    expect_equal(result$taxonomy_levels, c("Phylum", "Class", "Genus"))
    
    # Test 2: Only generates files for custom levels
    file_checks <- check_output_files("test_custom", c("Phylum", "Class", "Genus"))
    expect_true(all(file_checks),
                info = paste("Missing files:", 
                             paste(names(file_checks)[!file_checks], collapse = ", ")))
    
    # Test 3: Check summary file exists
    expect_true(file.exists(result$summary_file))
    
    # Clean up output files
    cleanup_output_files("test_custom", c("Phylum", "Class", "Genus"))
    
  }, finally = {
    # Clean up test files
    cleanup_test_files(test_files)
  })
})

test_that("mtap_analyze handles errors appropriately", {
  test_files <- create_test_files()
  
  tryCatch({
    # Test 1: Invalid config object
    invalid_config <- list(otu_file = "test.csv", tax_file = "test.csv")
    expect_error(
      mtap_analyze(invalid_config),
      "config must be a mtap_config object"
    )
    
    # Test 2: Non-existent file in config
    config_bad <- mtap_config(
      otu_file = "/nonexistent/file.csv",
      tax_file = test_files$tax_file,
      output_prefix = "test_error"
    )
    
    expect_error(
      suppressMessages(mtap_analyze(config_bad)),
      "OTU table file does not exist"
    )
    
  }, finally = {
    # Clean up test files
    cleanup_test_files(test_files)
  })
})

test_that("mtap_analyze with interactive mode", {
  test_files <- create_test_files()
  
  tryCatch({
    # Create configuration
    config <- mtap_config(
      otu_file = test_files$otu_file,
      tax_file = test_files$tax_file,
      output_prefix = "test_interactive"
    )
    
    # Capture output to check interactive messages
    output <- capture.output({
      result <- suppressMessages(mtap_analyze(config, interactive = TRUE))
    }, type = "output")
    
    # Test 1: Check for interactive-specific messages
    interactive_messages <- any(grepl("Analysis results saved to current working directory", output)) ||
      any(grepl("Use `print\\(result\\)` to view result summary", output))
    
    # Note: We don't require these messages to always appear, but if they do, they should be correct
    
    # Test 2: Result is still valid
    expect_s3_class(result, "mtap_result")
    
    # Clean up output files
    default_levels <- c("Kingdom", "Phylum", "Class", "Order", "Family", "Genus", "Species")
    cleanup_output_files("test_interactive", default_levels)
    
  }, finally = {
    # Clean up test files
    cleanup_test_files(test_files)
  })
})

# =============================================
# Tests for mtap_quick()
# =============================================

test_that("mtap_quick works like mtap_analyze", {
  test_files <- create_test_files()
  
  tryCatch({
    # Run quick analysis
    capture.output({
      result <- suppressMessages(
        mtap_quick(
          otu_file = test_files$otu_file,
          tax_file = test_files$tax_file,
          output_prefix = "test_quick"
        )
      )
    }, type = "output")
    
    # Test 1: Returns correct class
    expect_s3_class(result, "mtap_result")
    
    # Test 2: Contains expected components
    expected_components <- c(
      "config", "otu_data", "tax_data", "merged_data",
      "taxonomy_levels", "abundance_tables", "summary_file",
      "output_files"
    )
    expect_named(result, expected_components, ignore.order = TRUE)
    
    # Test 3: Config was created automatically
    expect_s3_class(result$config, "mtap_config")
    expect_equal(
      normalizePath(result$config$otu_file, mustWork = FALSE, winslash = "/"),
      normalizePath(test_files$otu_file, mustWork = FALSE, winslash = "/")
    )
    expect_equal(
      normalizePath(result$config$tax_file, mustWork = FALSE, winslash = "/"),
      normalizePath(test_files$tax_file, mustWork = FALSE, winslash = "/")
    )
    expect_equal(result$config$output_prefix, "test_quick")
    
    # Clean up output files
    default_levels <- c("Kingdom", "Phylum", "Class", "Order", "Family", "Genus", "Species")
    cleanup_output_files("test_quick", default_levels)
    
  }, finally = {
    # Clean up test files
    cleanup_test_files(test_files)
  })
})

test_that("mtap_quick passes additional parameters to mtap_config", {
  test_files <- create_test_files()
  
  tryCatch({
    # Run quick analysis with custom parameters
    capture.output({
      result <- suppressMessages(
        mtap_quick(
          otu_file = test_files$otu_file,
          tax_file = test_files$tax_file,
          output_prefix = "test_quick_params",
          taxonomy_hierarchy = c("Phylum", "Genus"),
          min_abundance = 0.01,
          min_prevalence = 0.1
        )
      )
    }, type = "output")
    
    # Test 1: Custom parameters were applied
    expect_equal(result$config$taxonomy_hierarchy, c("Phylum", "Genus"))
    expect_equal(result$config$min_abundance, 0.01)
    expect_equal(result$config$min_prevalence, 0.1)
    expect_equal(result$config$output_prefix, "test_quick_params")
    
    # Test 2: Uses custom taxonomy levels
    expect_equal(result$taxonomy_levels, c("Phylum", "Genus"))
    
    # Clean up output files
    cleanup_output_files("test_quick_params", c("Phylum", "Genus"))
    
  }, finally = {
    # Clean up test files
    cleanup_test_files(test_files)
  })
})

# =============================================
# Tests for print.mtap_result()
# =============================================

test_that("print.mtap_result works correctly", {
  test_files <- create_test_files()
  
  tryCatch({
    # Create configuration and run analysis
    config <- mtap_config(
      otu_file = test_files$otu_file,
      tax_file = test_files$tax_file,
      output_prefix = "test_print"
    )
    
    capture.output({
      result <- suppressMessages(mtap_analyze(config))
    }, type = "output")
    
    # Test 1: Print method returns the object invisibly
    printed_output <- capture.output(print_result <- print(result))
    expect_identical(print_result, result)
    
    # Test 2: Output contains expected information
    expect_true(any(grepl("MTaP Analysis Result Summary", printed_output)))
    expect_true(any(grepl("Analysis Configuration", printed_output)))
    expect_true(any(grepl("Data Statistics", printed_output)))
    expect_true(any(grepl("OTU count", printed_output)))
    expect_true(any(grepl("Sample count", printed_output)))
    expect_true(any(grepl("Taxonomy levels", printed_output)))
    expect_true(any(grepl("Generated Files", printed_output)))
    
    # Test 3: Files are listed
    default_levels <- c("Kingdom", "Phylum", "Class", "Order", "Family", "Genus", "Species")
    for (level in default_levels) {
      file_pattern <- paste0("test_print_", tolower(level), ".csv")
      expect_true(any(grepl(file_pattern, printed_output)) || 
                    !file.exists(paste0("test_print_", tolower(level), ".csv")))
    }
    
    # Clean up output files
    cleanup_output_files("test_print", default_levels)
    
  }, finally = {
    # Clean up test files
    cleanup_test_files(test_files)
  })
})

test_that("print.mtap_result handles missing files gracefully", {
  test_files <- create_test_files()
  
  tryCatch({
    # Create a mock result without actual file creation
    config <- mtap_config(
      otu_file = test_files$otu_file,
      tax_file = test_files$tax_file,
      output_prefix = "nonexistent_prefix"
    )
    
    # Mock result without running full analysis
    mock_result <- list(
      config = config,
      otu_data = data.frame(OTU_ID = "OTU_1", Sample1 = 100),
      tax_data = data.frame(OTU_ID = "OTU_1", Kingdom = "Bacteria"),
      merged_data = data.frame(OTU_ID = "OTU_1", Sample1 = 100, Kingdom = "Bacteria"),
      taxonomy_levels = c("Kingdom"),
      abundance_tables = list(Kingdom = data.frame(Kingdom = "Bacteria", Sample1 = 100)),
      summary_file = "nonexistent_prefix_summary.txt",
      output_files = list(
        abundance = c("nonexistent_prefix_kingdom.csv"),
        summary = "nonexistent_prefix_summary.txt"
      )
    )
    class(mock_result) <- c("mtap_result", "list")
    
    # Should not crash even if files don't exist
    printed_output <- capture.output(print(mock_result))
    
    # Should still display basic information
    expect_true(any(grepl("MTaP Analysis Result Summary", printed_output)))
    expect_true(any(grepl("Data Statistics", printed_output)))
    
  }, finally = {
    # Clean up test files
    cleanup_test_files(test_files)
  })
})

# =============================================
# Tests for mtap_help()
# =============================================

test_that("mtap_help displays help information", {
  # Capture help output
  help_output <- capture.output(mtap_help())
  
  # Test 1: Output contains package name
  expect_true(any(grepl("MTaP", help_output)))
  
  # Test 2: Output contains function categories
  expect_true(any(grepl("Main Functions", help_output)))
  expect_true(any(grepl("Data Reading Functions", help_output)))
  expect_true(any(grepl("Data Processing Functions", help_output)))
  expect_true(any(grepl("Utility Functions", help_output)))
  expect_true(any(grepl("Usage Examples", help_output)))
  
  # Test 3: Contains main function names
  expect_true(any(grepl("mtap_config", help_output)))
  expect_true(any(grepl("mtap_analyze", help_output)))
  expect_true(any(grepl("mtap_quick", help_output)))
  expect_true(any(grepl("mtap_cli", help_output)))
  
  # Test 4: Returns invisibly
  expect_invisible(mtap_help())
})

test_that("mtap_help displays help information", {
  # Test that function returns NULL (invisibly)
  result <- mtap_help()
  expect_null(result)
  #expect_invisible(result)
})

# =============================================
# Edge cases and integration tests
# =============================================

test_that("mtap_analyze handles edge cases", {
  test_files <- create_test_files()
  
  tryCatch({
    # Test with minimal data
    config_minimal <- mtap_config(
      otu_file = test_files$otu_file,
      tax_file = test_files$tax_file,
      output_prefix = "test_minimal",
      taxonomy_hierarchy = c("Kingdom")  # Only one level
    )
    
    capture.output({
      result <- suppressMessages(mtap_analyze(config_minimal))
    }, type = "output")
    
    # Should still work with single taxonomy level
    expect_s3_class(result, "mtap_result")
    expect_equal(result$taxonomy_levels, "Kingdom")
    
    # Clean up
    cleanup_output_files("test_minimal", "Kingdom")
    
  }, finally = {
    cleanup_test_files(test_files)
  })
})

test_that("mtap_analyze result structure is consistent", {
  test_files <- create_test_files()
  
  tryCatch({
    config <- mtap_config(
      otu_file = test_files$otu_file,
      tax_file = test_files$tax_file,
      output_prefix = "test_consistency"
    )
    
    capture.output({
      result <- suppressMessages(mtap_analyze(config))
    }, type = "output")
    
    # Test 1: Data frames have expected structure
    expect_true(ncol(result$otu_data) >= 1)
    expect_true(ncol(result$tax_data) >= 1)
    
    # Get first column names
    otu_first_col <- names(result$otu_data)[1]
    tax_first_col <- names(result$tax_data)[1]
    
    # Test 2: First columns should be the same (OTU identifier)
    expect_equal(otu_first_col, tax_first_col,
                 info = paste("First columns don't match:", 
                              otu_first_col, "vs", tax_first_col))
    
    # Test 3: Check that merged data contains taxonomy information
    taxonomy_cols <- c("Kingdom", "Phylum", "Class", "Order", "Family", "Genus", "Species")
    has_taxonomy <- any(taxonomy_cols %in% names(result$merged_data))
    expect_true(has_taxonomy,
                info = "Merged data doesn't contain taxonomy columns")
    
    # Test 4: Check sample data is preserved
    expect_true(ncol(result$merged_data) >= 2)
    
    # Test 5: Abundance tables list has correct names
    expect_named(result$abundance_tables, result$taxonomy_levels)
    
    # Clean up
    default_levels <- c("Kingdom", "Phylum", "Class", "Order", "Family", "Genus", "Species")
    cleanup_output_files("test_consistency", default_levels)
    
  }, finally = {
    cleanup_test_files(test_files)
  })
})