#include "base.h"
#include "platform.h"

// ----------------------------- Globals ------------------------------------
static atomic_uint_fast64_t g_files_found = 0;
static atomic_uint_fast64_t g_files_hashed = 0;
static atomic_uint_fast64_t g_bytes_processed = 0;
static atomic_int g_scan_done = 0;

// ============================= Utils ======================================
// ----------------------------- Timer functions --------------
static void timer_init(void) { QueryPerformanceFrequency(&g_qpc_freq); }

static void timer_start(HiResTimer *t) { QueryPerformanceCounter(&t->start); }

static double timer_stop(HiResTimer *t) {
  QueryPerformanceCounter(&t->end);
  return (double)(t->end.QuadPart - t->start.QuadPart) /
         (double)g_qpc_freq.QuadPart;
}

// ----------------------------- Get instruction set --------------
const char *get_xxhash_instruction_set(void) {
  int vecID = XXH_featureTest();

  switch (vecID) {
  case XXH_SCALAR:
    return "Scalar (portable C)";
  case XXH_SSE2:
    return "SSE2";
  case XXH_AVX2:
    return "AVX2";
  case XXH_AVX512:
    return "AVX-512";
  default:
    return "Unknown";
  }
}

// -------------------- Path parsing -------------------
static void normalize_path(char *p) {
  char *src = p;
  char *dst = p;
  int prev_slash = 0;

  while (*src) {
    char c = *src++;

    if (c == '\\' || c == '/') {
      if (!prev_slash) {
        *dst++ = '/';
        prev_slash = 1;
      }
    } else {
      *dst++ = c;
      prev_slash = 0;
    }
  }

  *dst = '\0';
}

static int parse_paths(char *line, char folders[][MAX_PATHLEN],
                       int max_folders) {
  int count = 0;
  char *p = line;

  while (*p && count < max_folders) {

    while (*p && isspace((unsigned char)*p))
      p++;

    if (!*p)
      break;

    char *start;
    char quote = 0;

    if (*p == '"' || *p == '\'') {
      quote = *p++;
      start = p;

      while (*p && *p != quote)
        p++;
    } else {
      start = p;

      while (*p && !isspace((unsigned char)*p))
        p++;
    }

    size_t len = p - start;
    if (len >= MAX_PATHLEN)
      len = MAX_PATHLEN - 1;

    memcpy(folders[count], start, len);
    folders[count][len] = 0;

    normalize_path(folders[count]);

    count++;

    if (quote && *p == quote)
      p++;
  }
  return count;
}

// ----------------------------- Convert filetime to epoch --------------
static uint64_t filetime_to_epoch(const FILETIME *ft) {
  ULARGE_INTEGER ull;
  ull.LowPart = ft->dwLowDateTime;
  ull.HighPart = ft->dwHighDateTime;

  // Windows epoch (1601) → Unix epoch (1970)
  return (ull.QuadPart - 116444736000000000ULL) / 10000000ULL;
}
// ----------------------------- Format time helper -------------------------
static void format_time(uint64_t t, char *out, size_t out_sz) {
  if (t == 0) {
    snprintf(out, out_sz, "N/A");
    return;
  }

  time_t tt = (time_t)t;
  struct tm tm;

#if PLATFORM_WINDOWS
  localtime_s(&tm, &tt);
#else
  localtime_r(&tt, &tm);
#endif

  strftime(out, out_sz, "%Y-%m-%d %H:%M:%S", &tm);
}

// ----------------------------- Resolve file owner ---------------------
static void get_file_owner(const char *path, char *out, size_t out_sz) {
  PSID sid = NULL;
  PSECURITY_DESCRIPTOR sd = NULL;

  if (GetNamedSecurityInfoA(path, SE_FILE_OBJECT, OWNER_SECURITY_INFORMATION,
                            &sid, NULL, NULL, NULL, &sd) == ERROR_SUCCESS) {

    char name[64], domain[64];
    DWORD name_len = sizeof(name);
    DWORD domain_len = sizeof(domain);
    SID_NAME_USE use;

    if (LookupAccountSidA(NULL, sid, name, &name_len, domain, &domain_len,
                          &use)) {
      snprintf(out, out_sz, "%s\\%s", domain, name);
    } else {
      snprintf(out, out_sz, "UNKNOWN");
    }
  } else {
    snprintf(out, out_sz, "UNKNOWN");
  }

  if (sd)
    LocalFree(sd);
}

// ----------------------------- Get file metadata -------------------------
void platform_get_file_times(const char *path, uint64_t *out_created,
                             uint64_t *out_modified) {
  WIN32_FILE_ATTRIBUTE_DATA fad;
  if (GetFileAttributesExA(path, GetFileExInfoStandard, &fad)) {
    *out_created = filetime_to_epoch(&fad.ftCreationTime);
    *out_modified = filetime_to_epoch(&fad.ftLastWriteTime);
  } else {
    *out_created = 0;
    *out_modified = 0;
  }
}

void platform_get_file_owner(const char *path, char *out_owner,
                             size_t out_owner_size) {
  get_file_owner(path, out_owner, out_owner_size);
}

// ----------------------------- Scan helpers -----------------------------
void scan_folder_windows_parallel(const char *base, ScannerContext *ctx) {
  char search[MAX_PATHLEN];
  snprintf(search, sizeof(search), "%s\\*", base);

  WIN32_FIND_DATAA fd;
  HANDLE h = FindFirstFileA(search, &fd);
  if (h == INVALID_HANDLE_VALUE)
    return;

  do {
    if (!strcmp(fd.cFileName, ".") || !strcmp(fd.cFileName, ".."))
      continue;

    char full[MAX_PATHLEN];
    snprintf(full, sizeof(full), "%s\\%s", base, fd.cFileName);

    if (fd.dwFileAttributes & FILE_ATTRIBUTE_REPARSE_POINT)
      continue;

    if (fd.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) {
      size_t len = strlen(full) + 1;

      char *dir = arena_push(&ctx->path_arena, len, false);
      memcpy(dir, full, len);

      mpmc_push_work(ctx->dir_queue, dir);
    } else {

      atomic_fetch_add(&g_files_found, 1);

      FileEntry *fe = arena_push(&ctx->meta_arena, sizeof(FileEntry), true);

      char norm[MAX_PATHLEN];
      strncpy(norm, full, sizeof(norm) - 1);
      norm[sizeof(norm) - 1] = 0;
      normalize_path(norm);

      size_t len = strlen(norm) + 1;

      char *path = arena_push(&ctx->path_arena, len, false);
      memcpy(path, norm, len);

      fe->path = path;

      platform_get_file_times(full, &fe->created_time, &fe->modified_time);

      platform_get_file_owner(full, fe->owner, sizeof(fe->owner));

      fe->size_bytes = ((uint64_t)fd.nFileSizeHigh << 32) | fd.nFileSizeLow;

      mpmc_push(ctx->file_queue, fe);
    }

  } while (FindNextFileA(h, &fd));

  FindClose(h);
}

// ------------------------- Scan worker --------------------------------
static DWORD WINAPI scan_worker(LPVOID arg) {
  ScannerContext *ctx = arg;

  for (;;) {
    char *dir = mpmc_pop(ctx->dir_queue);
    if (!dir)
      break;

    scan_folder_windows_parallel(dir, ctx);

    mpmc_task_done(ctx->dir_queue, ctx->num_threads);
  }
  return 0;
}

// ----------------------------- Hashing helpers -----------------------------
static void xxh3_hash_file_stream(const char *path, char *out_hex, BYTE *buf) {
  // compute XXH3_128 over file. POSIX and Windows use standard reads in this
  // helper.
  // On Windows try to use overlapped synchronous chunked reads for higher
  // throughput.
  HANDLE hFile =
      CreateFileA(path, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL,
                  OPEN_EXISTING, FILE_FLAG_SEQUENTIAL_SCAN, NULL);
  if (hFile == INVALID_HANDLE_VALUE) {
    strcpy(out_hex, "ERROR");
    return;
  }
  XXH128_hash_t h;
  XXH3_state_t state;
  XXH3_128bits_reset(&state);

  DWORD read = 0;
  BOOL ok;
  while (ReadFile(hFile, buf, READ_BLOCK, &read, NULL) && read > 0) {
    XXH3_128bits_update(&state, buf, (size_t)read);
    atomic_fetch_add(&g_bytes_processed, (uint64_t)read);
  }
  h = XXH3_128bits_digest(&state);
  CloseHandle(hFile);
  snprintf(out_hex, HASH_STRLEN, "%016llx%016llx", (unsigned long long)h.high64,
           (unsigned long long)h.low64);
}

// ------------------------- Hash worker --------------------------------
static DWORD WINAPI hash_worker(LPVOID arg) {

  WorkerContext *ctx = (WorkerContext *)arg;
  BYTE *buf = (BYTE *)malloc(READ_BLOCK);

  for (;;) {
    FileEntry *fe = mpmc_pop(ctx->file_queue);
    if (!fe)
      break;

    char hash[HASH_STRLEN];
    xxh3_hash_file_stream(fe->path, hash, buf);

    char created[32], modified[32];
    format_time(fe->created_time, created, sizeof(created));
    format_time(fe->modified_time, modified, sizeof(modified));

    double size_kib = (double)fe->size_bytes / 1024.0;

    char stack_buf[1024];

    int len =
        snprintf(stack_buf, sizeof(stack_buf), "%s\t%s\t%.2f\t%s\t%s\t%s\n",
                 hash, fe->path, size_kib, created, modified, fe->owner);

    char *dst = arena_push(&ctx->arena, len, false);
    memcpy(dst, stack_buf, len);

    atomic_fetch_add(&g_files_hashed, 1);
  }
  // free(buf); It will be freed by the system when the program exits

  return 0;
}

// ----------------------------- Progress display ---------------------------
DWORD WINAPI progress_thread(void *arg) {

  LARGE_INTEGER freq, start;
  QueryPerformanceFrequency(&freq);
  QueryPerformanceCounter(&start);

  uint64_t last_bytes = atomic_load(&g_bytes_processed);
  double last_time = 0.0;

  double displayed_speed = 0.0;
  const double sample_interval = 0.5;

  for (;;) {

    uint64_t found = atomic_load(&g_files_found);
    uint64_t hashed = atomic_load(&g_files_hashed);
    uint64_t bytes = atomic_load(&g_bytes_processed);
    int scan_done = atomic_load(&g_scan_done);

    LARGE_INTEGER now;
    QueryPerformanceCounter(&now);

    double t = (double)(now.QuadPart - start.QuadPart) / (double)freq.QuadPart;

    if (last_time == 0.0) {
      last_time = t;
      last_bytes = bytes;
    }

    double dt = t - last_time;

    if (dt >= sample_interval) {
      uint64_t db = bytes - last_bytes;

      if (db > 0 && dt > 0.0001) {
        displayed_speed = (double)db / (1024.0 * 1024.0) / dt;
      }

      last_bytes = bytes;
      last_time = t;
    }

    if (!scan_done) {

      printf("\rScanning: %llu files | Hashed: %llu | %.2f MB/s    ",
             (unsigned long long)found, (unsigned long long)hashed,
             displayed_speed);

    } else {

      double pct = found ? (double)hashed / (double)found : 0.0;

      int barw = 40;
      int filled = (int)(pct * barw);

      char bar[64];
      int p = 0;

      bar[p++] = '[';

      for (int i = 0; i < filled; i++)
        bar[p++] = '#';

      for (int i = filled; i < barw; i++)
        bar[p++] = '.';

      bar[p++] = ']';
      bar[p] = 0;

      printf("\r%s %6.2f%% (%llu / %llu) %.2f MB/s    ", bar, pct * 100.0,
             (unsigned long long)hashed, (unsigned long long)found,
             displayed_speed);
    }

    fflush(stdout);

    if (scan_done && hashed == found)
      break;

    Sleep(100);
  }

  printf("\n");

  return 0;
}

// ----------------------------- Main ---------------------------------------
int main(int argc, char **argv) {
  char folders[64][MAX_PATHLEN]; // up to 64 input folders
  int folder_count = 0;

  // -------------------------------
  // Parse arguments
  // -------------------------------
  for (int i = 1; i < argc; ++i) {
    if (folder_count < 64) {
      normalize_path(argv[i]);
      strncpy(folders[folder_count], argv[i], MAX_PATHLEN - 1);
      folders[folder_count][MAX_PATHLEN - 1] = 0;
      folder_count++;
    }
  }

  // -------------------------------
  // Ask user if no folders provided
  // -------------------------------
  if (folder_count == 0) {
    printf("Enter folders to process (Enter = current folder): ");
    fflush(stdout);

    char buf[KiB(32)];

    if (!fgets(buf, sizeof(buf), stdin))
      return 1;

    buf[strcspn(buf, "\r\n")] = 0;

    if (buf[0] == 0) {
      strcpy(folders[0], ".");
      folder_count = 1;
    } else {
      folder_count = parse_paths(buf, folders, 64);
    }
  }

  // Display selected folders
  printf("Processing %d folder(s):\n", folder_count);
  for (int i = 0; i < folder_count; ++i) {
    printf("  - %s\n", folders[i]);
  }

  // -------------------------------
  // Scanning and total timer init
  // -------------------------------
  timer_init();

  HiResTimer total_timer;
  HiResTimer scan_timer;

  timer_start(&total_timer);
  timer_start(&scan_timer);

  // -------------------------------
  // Creating a general purpose arena
  // -------------------------------
  arena_params params = {
      .reserve_size = GiB(1),
      .commit_size = MiB(16),
      .align = 0,
      .push_size = 0,
      .allow_free_list = true,
      .allow_swapback = false,
      .growth_policy = ARENA_GROWTH_NORMAL,
      .commit_policy = ARENA_COMMIT_LAZY,
      .max_nbre_blocks = 1,
  };

  mem_arena *gp_arena = arena_create(&params);

  // -------------------------------
  // Detect hardware threads
  // -------------------------------
  size_t hw_threads = 1;
  // --- Windows: detect PHYSICAL cores (not logical threads) ---
  DWORD len = 0;
  GetLogicalProcessorInformation(NULL, &len);

  SYSTEM_LOGICAL_PROCESSOR_INFORMATION *buf =
      (SYSTEM_LOGICAL_PROCESSOR_INFORMATION *)arena_push(&gp_arena, len, true);

  if (GetLogicalProcessorInformation(buf, &len)) {
    DWORD count = 0;
    DWORD n = len / sizeof(SYSTEM_LOGICAL_PROCESSOR_INFORMATION);
    for (DWORD i = 0; i < n; i++) {
      if (buf[i].Relationship == RelationProcessorCore)
        count++;
    }
    if (count > 0)
      hw_threads = count;
  }
  arena_free(&gp_arena, (u8 **)&buf, len);

  // Logical threads = CPU cores * 2
  size_t num_threads = hw_threads * 2;

  printf("Starting thread pool: %zu threads (CPU cores: %zu)\n", num_threads,
         hw_threads);
  printf("  Selected instruction set: %s\n", get_xxhash_instruction_set());

  // -------------------------------
  // Scanning and hashing
  // -------------------------------
  MPMCQueue dir_queue;
  mpmc_init(&dir_queue, MiB(1));

  MPMCQueue file_queue;
  mpmc_init(&file_queue, MiB(1));

  // starting hash threads
  size_t num_hash_threads = num_threads;

  WorkerContext workers[num_hash_threads];

  HANDLE *hash_threads =
      arena_push(&gp_arena, sizeof(HANDLE) * num_hash_threads, true);

  for (size_t i = 0; i < num_hash_threads; ++i) {

    workers[i].arena = arena_create(&params);
    workers[i].file_queue = &file_queue;

    hash_threads[i] = CreateThread(NULL, 0, hash_worker, &workers[i], 0, NULL);
  }

  // starting progress printing thread
  HANDLE progress = CreateThread(NULL, 0, progress_thread, NULL, 0, NULL);

  // starting scan threads
  size_t num_scan_threads = num_threads;

  ScannerContext scanners[num_scan_threads];

  HANDLE *scan_threads =
      arena_push(&gp_arena, sizeof(HANDLE) * num_scan_threads, true);

  for (size_t i = 0; i < num_scan_threads; i++) {
    scanners[i].num_threads = num_scan_threads;

    scanners[i].path_arena = arena_create(&params);
    scanners[i].meta_arena = arena_create(&params);

    scanners[i].dir_queue = &dir_queue;
    scanners[i].file_queue = &file_queue;

    scan_threads[i] = CreateThread(NULL, 0, scan_worker, &scanners[i], 0, NULL);
  }

  // Initial folder push
  for (int i = 0; i < folder_count; i++) {

    size_t len = strlen(folders[i]) + 1;

    char *path = arena_push(&scanners[0].path_arena, len, false);
    memcpy(path, folders[i], len);

    mpmc_push_work(&dir_queue, path);
  }

  // Stop scan threads
  WaitForMultipleObjects((DWORD)num_scan_threads, scan_threads, TRUE, INFINITE);

  for (size_t i = 0; i < num_scan_threads; ++i)
    CloseHandle(scan_threads[i]);

  mpmc_producers_finished(&file_queue, num_hash_threads);

  atomic_store(&g_scan_done, 1);

  arena_free(&gp_arena, (u8 **)&scan_threads,
             sizeof(HANDLE) * num_scan_threads);

  double scan_seconds = timer_stop(&scan_timer);
  size_t total_found = atomic_load(&g_files_found);

  printf("\r%*s\r", 120, ""); // clear_console_line
  printf("Completed scanning in %.2f seconds, found %zu files\n\n",
         scan_seconds, total_found);

  // If no files found
  if (total_found == 0) {
    printf("No files found.\n");
    return 0;
  }

  // Stop hashing threads
  WaitForMultipleObjects((DWORD)num_hash_threads, hash_threads, TRUE, INFINITE);

  for (size_t i = 0; i < num_hash_threads; ++i)
    CloseHandle(hash_threads[i]);

  arena_free(&gp_arena, (u8 **)&hash_threads,
             sizeof(HANDLE) * num_hash_threads);

  // Stop progress printing thread
  WaitForSingleObject(progress, INFINITE);
  CloseHandle(progress);

  // -------------------------------
  // Export file_hashes.txt
  // -------------------------------

  FILE *f = fopen(FILE_HASHES_TXT, "wb");

  for (int i = 0; i < num_threads; i++) {
    mem_arena *arena = workers[i].arena;

    u8 *arena_base =
        (u8 *)arena + ALIGN_UP_POW2(sizeof(mem_arena), arena->align);
    fwrite(arena_base, 1, arena->pos, f);
  }

  fclose(f);

  // -------------------------------
  // Print summary
  // -------------------------------
  double total_seconds = timer_stop(&total_timer);

  printf("Completed hashing %zu files\n", total_found);

  uint64_t total_bytes = (uint64_t)atomic_load(&g_bytes_processed);
  double total_mb = (double)total_bytes / (1024.0 * 1024.0);
  double avg_mbps = total_mb / total_seconds;
  printf("Total: %.2f MB, Average: %.2f MB/s\n", total_mb, avg_mbps);
  printf("  Total time  : %.2f seconds\n\n", total_seconds);

  return 0;
}
