#include <cstdio>
#include <cstdint>
#include <cstring>
#include <string>
#include <vector>
#include <algorithm>
#include <filesystem>
#include <glob.h>

namespace fs = std::filesystem;

static constexpr double CELL = 0.5;
static constexpr int SUB = 32;
static constexpr double LON_MIN = -180.0;
static constexpr double LAT_MIN = -60.0;
static constexpr double LON_MAX = 180.0;
static constexpr double LAT_MAX = 85.0;

struct RawPoint {
    int32_t col;
    int32_t row;
    uint16_t lonq;
    uint16_t latq;
    uint8_t speed;
};

struct Point64 {
    uint16_t lonq;
    uint16_t latq;
    uint8_t speed;
};

static bool cmp_cell(const RawPoint& a, const RawPoint& b) {
    if (a.col != b.col) return a.col < b.col;
    if (a.row != b.row) return a.row < b.row;
    return false;
}

int main(int argc, char* argv[]) {
    if (argc < 3) {
        fprintf(stderr, "Usage: %s <output_dir> <raw_file1> [raw_file2] ...\n", argv[0]);
        return 1;
    }

    const char* outdir = argv[1];
    fs::create_directories(outdir);

    std::vector<RawPoint> all_points;

    for (int i = 2; i < argc; i++) {
        FILE* f = fopen(argv[i], "rb");
        if (!f) { fprintf(stderr, "Cannot open %s\n", argv[i]); continue; }

        fseek(f, 0, SEEK_END);
        long fsize = ftell(f);
        fseek(f, 0, SEEK_SET);

        uint64_t before = all_points.size();
        uint64_t npoints = fsize / sizeof(RawPoint);
        all_points.resize(before + npoints);
        fread(all_points.data() + before, sizeof(RawPoint), npoints, f);
        fclose(f);

        fprintf(stderr, "Loaded %s: %lu points (total: %lu)\n", argv[i], npoints, all_points.size());
    }

    fprintf(stderr, "Sorting %lu points...\n", all_points.size());
    std::sort(all_points.begin(), all_points.end(), cmp_cell);
    fprintf(stderr, "Sort done. Writing cells...\n");

    std::string index_path = std::string(outdir) + "/index.bin";
    FILE* index = fopen(index_path.c_str(), "wb");
    if (!index) { fprintf(stderr, "Cannot write index\n"); return 1; }

    uint32_t magic = 0x5A4D4150;
    fwrite(&magic, 4, 1, index);
    uint16_t version = 2;
    fwrite(&version, 2, 1, index);
    uint16_t sub = SUB;
    fwrite(&sub, 2, 1, index);
    fwrite(&CELL, 8, 1, index);
    fwrite(&LON_MIN, 8, 1, index);
    fwrite(&LAT_MIN, 8, 1, index);
    fwrite(&LON_MAX, 8, 1, index);
    fwrite(&LAT_MAX, 8, 1, index);

    uint32_t num_cells = 0;
    fseek(index, 36, SEEK_SET);
    fwrite(&num_cells, 4, 1, index);
    fseek(index, 0, SEEK_END);

    int32_t cur_col = -1, cur_row = -1;
    FILE* cur_file = nullptr;
    uint32_t cell_points = 0;
    std::string cur_key;

    for (uint64_t i = 0; i <= all_points.size(); i++) {
        bool new_cell = (i == all_points.size() ||
                         all_points[i].col != cur_col ||
                         all_points[i].row != cur_row);

        if (new_cell && cur_file) {
            fclose(cur_file);
            uint32_t klen = static_cast<uint32_t>(cur_key.size());
            fwrite(&klen, 4, 1, index);
            fwrite(cur_key.data(), 1, klen, index);
            fwrite(&cell_points, 4, 1, index);
            num_cells++;
            if (num_cells % 500 == 0) {
                fprintf(stderr, "  %u cells written...\n", num_cells);
            }
        }

        if (i == all_points.size()) break;

        if (new_cell) {
            cur_col = all_points[i].col;
            cur_row = all_points[i].row;
            cur_key = std::to_string(cur_col) + "_" + std::to_string(cur_row);
            std::string cellfile = std::string(outdir) + "/cell_" + cur_key + ".bin";
            cur_file = fopen(cellfile.c_str(), "wb");
            cell_points = 0;
        }

        Point64 p;
        p.lonq = all_points[i].lonq;
        p.latq = all_points[i].latq;
        p.speed = all_points[i].speed;
        fwrite(&p, sizeof(Point64), 1, cur_file);
        cell_points++;
    }

    fseek(index, 36, SEEK_SET);
    fwrite(&num_cells, 4, 1, index);
    fclose(index);

    fprintf(stderr, "Total cells: %u\n", num_cells);
    fprintf(stderr, "Grid written to %s\n", outdir);
    return 0;
}