/*
# Compile
gcc -o io_uring_test io_uring_test2.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 BUFFER_SIZE 4096
#define NUM_BUFFERS 4
#define NUM_REGISTERED_FILES 3 // Test with 3 files

// 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);
}

static int create_test_file(const char *filename, const char *content) {
  FILE *f = fopen(filename, "w");
  if (!f) {
    perror("Failed to create test file");
    return -1;
  }

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

  printf("  Created test file: %s\n", filename);
  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");

  size_t total_size = BUFFER_SIZE * NUM_BUFFERS;
  *buffers = aligned_alloc(4096, total_size);
  if (!*buffers) {
    print_failure("Buffer allocation", strerror(errno));
    results->failed++;
    return -1;
  }

  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: Register files sparse (empty table)
static int test_register_files_sparse(struct io_uring *ring, unsigned nr_files,
                                      TestResults *results) {
  print_step("Sparse file registration (empty table)");

  int ret = io_uring_register_files_sparse(ring, nr_files);
  if (ret < 0) {
    if (ret == -EINVAL) {
      print_info(
          "io_uring_register_files_sparse not supported (kernel < 5.19)");
      print_info("Trying regular file registration with invalid fds...");

      // Fallback: register with invalid fds
      int *invalid_fds = calloc(nr_files, sizeof(int));
      if (!invalid_fds) {
        print_failure("Allocating invalid fds array", "Out of memory");
        results->failed++;
        return -1;
      }

      for (int i = 0; i < nr_files; i++) {
        invalid_fds[i] = -1; // Mark all as invalid
      }

      ret = io_uring_register_files(ring, invalid_fds, nr_files);
      free(invalid_fds);

      if (ret < 0) {
        print_failure("Regular file registration also failed", strerror(-ret));
        results->failed++;
        return -1;
      }
      print_success("File table registered (regular, with invalid fds)");
    } else {
      print_failure("io_uring_register_files_sparse", strerror(-ret));
      results->failed++;
      return -1;
    }
  } else {
    printf("  Registered empty file table with %u slots\n", nr_files);
    print_success("Sparse file table created");
  }

  results->passed++;
  return 0;
}

// Test 4: Update file slot and read from it
static int test_file_read_loop(struct io_uring *ring, struct iovec *iovs,
                               const char **filenames,
                               const char **expected_contents, int num_files,
                               TestResults *results) {
  print_step("File slot update and read loop");

  int *fds = calloc(num_files, sizeof(int));
  if (!fds) {
    print_failure("Allocating fd array", "Out of memory");
    results->failed++;
    return -1;
  }

  // Open all files first
  for (int i = 0; i < num_files; i++) {
    fds[i] = open(filenames[i], O_RDONLY);
    if (fds[i] < 0) {
      print_failure("Opening file", filenames[i]);
      results->failed++;
      // Close already opened files
      for (int j = 0; j < i; j++)
        close(fds[j]);
      free(fds);
      return -1;
    }
    printf("  Opened %s (fd=%d)\n", filenames[i], fds[i]);
  }

  // Test loop: update slot, submit read, verify
  for (int slot = 0; slot < num_files; slot++) {
    printf("\n  --- Testing slot %d with file '%s' ---\n", slot,
           filenames[slot]);

    // Update the file registration for this slot
    printf("  Updating slot %d with fd %d...\n", slot, fds[slot]);
    int ret = io_uring_register_files_update(ring, slot, &fds[slot], 1);

    if (ret < 0) {
      print_failure("File registration update", strerror(-ret));
      results->failed++;
      continue;
    }
    printf("  Slot update result: %d (expected 1)\n", ret);

    // Get file size for read size calculation
    struct stat st;
    if (fstat(fds[slot], &st) != 0) {
      print_failure("fstat", strerror(errno));
      results->failed++;
      continue;
    }

    size_t file_size = st.st_size;
    size_t read_size = BUFFER_SIZE;

    // Adjust read size for O_DIRECT if needed
    int page_size = plat_get_pagesize();
    if (read_size > file_size) {
      read_size = ALIGN_UP_POW2(file_size, page_size);
    }

    printf("  File size: %zu, read size: %zu\n", file_size, read_size);

    // Clear buffer for this test
    memset(iovs[0].iov_base, 0, BUFFER_SIZE);

    // Submit read using registered file
    struct io_uring_sqe *sqe = io_uring_get_sqe(ring);
    if (!sqe) {
      print_failure("Getting SQE", "No available SQE");
      results->failed++;
      continue;
    }

    // Use slot index with fixed file flag
    io_uring_prep_read_fixed(sqe, slot, iovs[0].iov_base, read_size, 0, 0);
    sqe->flags |= IOSQE_FIXED_FILE;
    io_uring_sqe_set_data64(sqe, 100 + slot); // Unique user_data per slot

    ret = io_uring_submit(ring);
    if (ret < 0) {
      print_failure("Submitting read", strerror(-ret));
      results->failed++;
      continue;
    }
    printf("  Submitted read (1 SQE)\n");

    // Wait for completion
    struct io_uring_cqe *cqe;
    ret = io_uring_wait_cqe(ring, &cqe);
    if (ret < 0) {
      print_failure("Waiting for completion", strerror(-ret));
      results->failed++;
      continue;
    }

    // Process completion
    uint64_t user_data = io_uring_cqe_get_data64(cqe);
    int bytes_read = cqe->res;

    printf("  Completion: user_data=%lu, result=%d\n", (unsigned long)user_data,
           bytes_read);

    if (bytes_read < 0) {
      print_failure("Read operation", strerror(-bytes_read));
      results->failed++;
      io_uring_cqe_seen(ring, cqe);
      continue;
    }

    if (user_data != 100 + slot) {
      print_failure("User data mismatch", "Wrong user_data value");
      results->failed++;
      io_uring_cqe_seen(ring, cqe);
      continue;
    }

    // Verify the data
    char *data = (char *)iovs[0].iov_base;
    printf("  Data read (%d bytes): %.*s\n", bytes_read,
           bytes_read < 100 ? bytes_read : 100, data);

    if (strstr(data, expected_contents[slot]) == NULL) {
      print_failure("Data verification",
                    "Expected content not found in read data");
      results->failed++;
    } else {
      print_success("Data verified successfully");
      results->passed++;
    }

    io_uring_cqe_seen(ring, cqe);

    // Invalidate the slot after use (mark as -1)
    int invalid_fd = -1;
    ret = io_uring_register_files_update(ring, slot, &invalid_fd, 1);
    if (ret < 0) {
      printf("  Warning: Could not invalidate slot %d: %s\n", slot,
             strerror(-ret));
    }
  }

  // Close all files
  for (int i = 0; i < num_files; i++) {
    if (fds[i] >= 0)
      close(fds[i]);
  }
  free(fds);

  return 0;
}

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

  printf(COLOR_BLUE "\n========================================\n");
  printf("    io_uring Sparse File Registration Test\n");
  printf("========================================\n" COLOR_RESET);

  // Define test files and their content
  const char *filenames[] = {"test_file_0.txt", "test_file_1.txt",
                             "test_file_2.txt"};

  const char *contents[] = {
      "This is file 0: Hello World! The quick brown fox jumps over the lazy "
      "dog.",
      "This is file 1: io_uring is awesome for async I/O operations!",
      "This is file 2: Testing sparse file registration with multiple files."};

  const char *expected_substrings[] = {"Hello World", "io_uring is awesome",
                                       "sparse file registration"};

  int num_files = 3;

  // Create all test files
  print_info("Creating test files...");
  for (int i = 0; i < num_files; i++) {
    if (create_test_file(filenames[i], contents[i]) != 0) {
      return 1;
    }
  }

  // Test 1: Create io_uring
  if (test_io_uring_create(&ring, &results) != 0) {
    goto cleanup_files;
  }

  // Test 2: Register buffers
  if (test_register_buffers(&ring, &buffers, iovs, &results) != 0) {
    io_uring_queue_exit(&ring);
    goto cleanup_files;
  }

  // Test 3: Register empty file table (sparse)
  if (test_register_files_sparse(&ring, num_files, &results) != 0) {
    io_uring_unregister_buffers(&ring);
    free(buffers);
    io_uring_queue_exit(&ring);
    goto cleanup_files;
  }

  // Test 4: Loop through files, update slots, read and verify
  test_file_read_loop(&ring, iovs, filenames, expected_substrings, num_files,
                      &results);

  // Cleanup
  io_uring_unregister_files(&ring);
  io_uring_unregister_buffers(&ring);
  free(buffers);
  io_uring_queue_exit(&ring);

cleanup_files:
  // Remove test files
  for (int i = 0; i < num_files; i++) {
    remove(filenames[i]);
  }

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

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