/*
# Compile
gcc -o io_uring_test io_uring_test.c -luring

# Run
./io_uring_test
   */
#include "base.h"
#include <stdint.h>
#define _GNU_SOURCE
#include <errno.h>
#include <fcntl.h>
#include <liburing.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <unistd.h>

#define TEST_FILE "test_io_uring.txt"
#define BUFFER_SIZE 4096
#define NUM_BUFFERS 4

// Colors for output
#define COLOR_GREEN "\033[0;32m"
#define COLOR_RED "\033[0;31m"
#define COLOR_YELLOW "\033[0;33m"
#define COLOR_BLUE "\033[0;34m"
#define COLOR_RESET "\033[0m"

// Test result tracking
typedef struct {
  int passed;
  int failed;
} TestResults;

static void print_success(const char *step) {
  printf(COLOR_GREEN "[✓] SUCCESS: %s" COLOR_RESET "\n", step);
}

static void print_failure(const char *step, const char *error) {
  printf(COLOR_RED "[✗] FAILED: %s - %s" COLOR_RESET "\n", step, error);
}

static void print_info(const char *msg) {
  printf(COLOR_BLUE "[i] INFO: %s" COLOR_RESET "\n", msg);
}

static void print_step(const char *step) {
  printf(COLOR_YELLOW "\n>>> Testing: %s" COLOR_RESET "\n", step);
}

// Create a test file with known content
static int create_test_file(void) {
  const char *test_content =
      "Hello, io_uring! This is a test file for async I/O operations.\n"
      "Line 2: Testing reads with registered buffers.\n"
      "Line 3: The quick brown fox jumps over the lazy dog.\n"
      "Line 4: ABCDEFGHIJKLMNOPQRSTUVWXYZ\n"
      "Line 5: 0123456789\n";

  FILE *f = fopen(TEST_FILE, "w");
  if (!f) {
    perror("Failed to create test file");
    return -1;
  }

  fprintf(f, "%s", test_content);
  fclose(f);

  print_info("Test file created successfully");
  return 0;
}

// Test 1: Create io_uring instance
static int test_io_uring_create(struct io_uring *ring, TestResults *results) {
  print_step("io_uring creation");

  int ret = io_uring_queue_init(256, ring, 0);
  if (ret < 0) {
    print_failure("io_uring_queue_init", strerror(-ret));
    results->failed++;
    return -1;
  }

  print_success("io_uring instance created");
  results->passed++;
  return 0;
}

// Test 2: Register buffers
static int test_register_buffers(struct io_uring *ring, void **buffers,
                                 struct iovec *iovs, TestResults *results) {
  print_step("Buffer registration");

  // Allocate and prepare buffers
  size_t total_size = BUFFER_SIZE * NUM_BUFFERS;
  *buffers = aligned_alloc(4096, total_size); // Page-aligned for O_DIRECT
  if (!*buffers) {
    print_failure("Buffer allocation", strerror(errno));
    results->failed++;
    return -1;
  }

  // Initialize iovecs
  for (int i = 0; i < NUM_BUFFERS; i++) {
    iovs[i].iov_base = (char *)*buffers + (i * BUFFER_SIZE);
    iovs[i].iov_len = BUFFER_SIZE;
    memset(iovs[i].iov_base, 0, BUFFER_SIZE);
  }

  int ret = io_uring_register_buffers(ring, iovs, NUM_BUFFERS);
  if (ret < 0) {
    print_failure("io_uring_register_buffers", strerror(-ret));
    results->failed++;
    return -1;
  }

  print_success("Buffers registered successfully");
  results->passed++;
  return 0;
}

// Test 3: Open file
// Modified test_open_file function
static int test_open_file(int *fd, TestResults *results) {
  print_step("File opening");

  // Get file size
  struct stat st;
  if (stat(TEST_FILE, &st) != 0) {
    print_failure("stat", strerror(errno));
    results->failed++;
    return -1;
  }

  // Check if file size is page-aligned
  int page_size = plat_get_pagesize();
  size_t file_size = st.st_size;

  printf("  File size: %zu bytes\n", file_size);
  printf("  Page size: %d bytes\n", page_size);

  if (file_size % page_size != 0) {
    printf("  Extending read size from %zu to %zu bytes\n", file_size,
           ALIGN_UP_POW2(file_size, page_size));
  }

  // Try O_DIRECT first
  *fd = open(TEST_FILE, O_RDONLY | O_DIRECT);
  if (*fd < 0) {
    print_info("O_DIRECT failed, trying without it");
    *fd = open(TEST_FILE, O_RDONLY);
    if (*fd < 0) {
      print_failure("open", strerror(errno));
      results->failed++;
      return -1;
    }
    print_info("Using buffered I/O (O_DIRECT not available)");
  } else {
    print_success("File opened with O_DIRECT");
  }

  results->passed++;
  return 0;
}

// Test 4: Build and submit read operation
static int test_submit_read(struct io_uring *ring, int fd, struct iovec *iovs,
                            int buffer_id, uint64_t user_data,
                            TestResults *results) {
  print_step("Building and submitting read operation");

  // Get file size for proper alignment
  struct stat st;
  if (fstat(fd, &st) != 0) {
    print_failure("fstat", strerror(errno));
    results->failed++;
    return -1;
  }

  u32 page_size = plat_get_pagesize();
  size_t file_size = st.st_size;
  size_t read_size = BUFFER_SIZE;

  // For O_DIRECT, ensure read size is sector-aligned
  if (read_size > file_size) {
    read_size = ALIGN_UP_POW2(file_size, page_size);
    printf("  Adjusted read size to %zu bytes for O_DIRECT alignment\n",
           read_size);
  }

  struct io_uring_sqe *sqe = io_uring_get_sqe(ring);
  if (!sqe) {
    print_failure("io_uring_get_sqe", "No available SQE");
    results->failed++;
    return -1;
  }

  // Prepare read operation using registered buffer
  io_uring_prep_read_fixed(sqe, fd, iovs[buffer_id].iov_base, read_size, 0,
                           buffer_id);
  io_uring_sqe_set_data64(sqe, user_data);

  int ret = io_uring_submit(ring);
  if (ret < 0) {
    print_failure("io_uring_submit", strerror(-ret));
    results->failed++;
    return -1;
  }

  print_success("Read operation submitted successfully");
  results->passed++;
  return 0;
}

// Test 5: Wait for completion
static int test_wait_completion(struct io_uring *ring,
                                struct io_uring_cqe **cqe,
                                TestResults *results) {
  print_step("Waiting for completion");

  int ret = io_uring_wait_cqe(ring, cqe);
  if (ret < 0) {
    print_failure("io_uring_wait_cqe", strerror(-ret));
    results->failed++;
    return -1;
  }

  print_success("Completion received");
  results->passed++;
  return 0;
}

// Test 6: Process completion
static int test_process_completion(struct io_uring_cqe *cqe,
                                   uint64_t expected_user_data,
                                   TestResults *results) {
  print_step("Processing completion");

  uint64_t user_data = io_uring_cqe_get_data64(cqe);
  int res = cqe->res;

  printf("  Completion data:\n");
  printf("    User data: %lu (expected: %lu)\n", user_data, expected_user_data);
  printf("    Result: %d bytes read\n", res);

  if (user_data != expected_user_data) {
    print_failure("User data mismatch",
                  "User data doesn't match expected value");
    results->failed++;
    return -1;
  }

  if (res < 0) {
    print_failure("Read operation", strerror(-res));
    results->failed++;
    return -1;
  }

  print_success("Completion processed successfully");
  results->passed++;
  return res; // Return number of bytes read
}

// Test 7: Verify read data
static int test_verify_data(struct iovec *iovs, int buffer_id, int bytes_read,
                            TestResults *results) {
  print_step("Data verification");

  char *data = (char *)iovs[buffer_id].iov_base;

  printf("  Read data (first 200 chars):\n");
  printf("  ---\n");
  for (int i = 0; i < bytes_read && i < 200; i++) {
    putchar(data[i]);
  }
  if (bytes_read > 200)
    printf("...");
  printf("\n  ---\n");

  // Check if data is not empty
  if (bytes_read == 0) {
    print_failure("Data verification", "No data read");
    results->failed++;
    return -1;
  }

  // Check if data contains expected content
  if (strstr(data, "io_uring") == NULL) {
    print_failure("Data verification", "Expected content not found");
    results->failed++;
    return -1;
  }

  print_success("Data verified successfully");
  results->passed++;
  return 0;
}

// Test 8: Test multiple concurrent reads
static int test_concurrent_reads(struct io_uring *ring, int fd,
                                 struct iovec *iovs, TestResults *results) {
  print_step("Concurrent reads test");

  int num_reads = 3;
  int submitted = 0;

  // Submit multiple reads
  for (int i = 0; i < num_reads; i++) {
    struct io_uring_sqe *sqe = io_uring_get_sqe(ring);
    if (!sqe) {
      print_failure("Getting SQE for concurrent read", "No available SQE");
      results->failed++;
      return -1;
    }

    off_t offset = i * 100; // Read from different offsets
    io_uring_prep_read_fixed(sqe, fd, iovs[i].iov_base, BUFFER_SIZE, offset, i);
    io_uring_sqe_set_data64(sqe, i);
    submitted++;
  }

  int ret = io_uring_submit(ring);
  if (ret != submitted) {
    char msg[64];
    snprintf(msg, sizeof(msg), "Expected %d, got %d", submitted, ret);

    print_failure("Submitting concurrent reads", msg);
    results->failed++;
    return -1;
  }

  print_success("Concurrent reads submitted");

  // Wait for and process completions
  for (int i = 0; i < submitted; i++) {
    struct io_uring_cqe *cqe;
    ret = io_uring_wait_cqe(ring, &cqe);
    if (ret < 0) {
      print_failure("Waiting for concurrent read completion", strerror(-ret));
      results->failed++;
      return -1;
    }

    uint64_t user_data = io_uring_cqe_get_data64(cqe);
    int res = cqe->res;

    printf("  Concurrent read %lu completed: %d bytes read\n", user_data, res);
    io_uring_cqe_seen(ring, cqe);
  }

  print_success("Concurrent reads completed successfully");
  results->passed++;
  return 0;
}

// Cleanup function
static void cleanup(struct io_uring *ring, int fd, void *buffers) {
  if (fd >= 0)
    close(fd);
  if (buffers) {
    io_uring_unregister_buffers(ring);
    free(buffers);
  }
  io_uring_queue_exit(ring);
  remove(TEST_FILE);
}

int main() {
  TestResults results = {0, 0};
  struct io_uring ring;
  int fd = -1;
  void *buffers = NULL;
  struct iovec iovs[NUM_BUFFERS];

  printf(COLOR_BLUE "\n========================================\n");
  printf("    io_uring Test Suite\n");
  printf("========================================\n" COLOR_RESET);

  // Create test file
  if (create_test_file() != 0) {
    return 1;
  }

  // Test 1: Create io_uring
  if (test_io_uring_create(&ring, &results) != 0) {
    cleanup(&ring, fd, buffers);
    return 1;
  }

  // Test 2: Register buffers
  if (test_register_buffers(&ring, &buffers, iovs, &results) != 0) {
    cleanup(&ring, fd, buffers);
    return 1;
  }

  // Test 3: Open file
  if (test_open_file(&fd, &results) != 0) {
    cleanup(&ring, fd, buffers);
    return 1;
  }

  // Test 4: Submit read
  uint64_t test_user_data = 12345;
  if (test_submit_read(&ring, fd, iovs, 0, test_user_data, &results) != 0) {
    cleanup(&ring, fd, buffers);
    return 1;
  }

  // Test 5: Wait for completion
  struct io_uring_cqe *cqe;
  if (test_wait_completion(&ring, &cqe, &results) != 0) {
    cleanup(&ring, fd, buffers);
    return 1;
  }

  // Test 6: Process completion
  int bytes_read = test_process_completion(cqe, test_user_data, &results);
  if (bytes_read < 0) {
    cleanup(&ring, fd, buffers);
    return 1;
  }
  io_uring_cqe_seen(&ring, cqe);

  // Test 7: Verify data
  if (test_verify_data(iovs, 0, bytes_read, &results) != 0) {
    cleanup(&ring, fd, buffers);
    return 1;
  }

  // Test 8: Concurrent reads
  if (test_concurrent_reads(&ring, fd, iovs, &results) != 0) {
    cleanup(&ring, fd, buffers);
    return 1;
  }

  // Print summary
  printf(COLOR_BLUE "\n========================================\n");
  printf("           TEST SUMMARY\n");
  printf("========================================\n" COLOR_RESET);
  printf("  Total tests: %d\n", results.passed + results.failed);
  printf(COLOR_GREEN "  Passed: %d\n" COLOR_RESET, results.passed);
  if (results.failed > 0) {
    printf(COLOR_RED "  Failed: %d\n" COLOR_RESET, results.failed);
  } else {
    printf(COLOR_GREEN "  ✓ ALL TESTS PASSED!\n" COLOR_RESET);
  }

  // Cleanup
  cleanup(&ring, fd, buffers);

  return results.failed > 0 ? 1 : 0;
}