#include <osmium/io/any_input.hpp>
#include <osmium/visitor.hpp>
#include <osmium/handler/node_locations_for_ways.hpp>
#include <osmium/index/map/sparse_file_array.hpp>
#include <osmium/osm/way.hpp>

#include <cstdio>
#include <cstdint>
#include <cstring>
#include <cmath>
#include <string>
#include <vector>
#include <algorithm>
#include <filesystem>
#include <fcntl.h>
#include <unistd.h>

namespace fs = std::filesystem;

using index_type = osmium::index::map::SparseFileArray<osmium::unsigned_object_id_type, osmium::Location>;
using location_handler_type = osmium::handler::NodeLocationsForWays<index_type>;

static constexpr double CELL = 0.5;
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;
};

static inline uint16_t quantize_lon(double lon) {
    return static_cast<uint16_t>((lon - LON_MIN) / (LON_MAX - LON_MIN) * 65535.0);
}

static inline uint16_t quantize_lat(double lat) {
    return static_cast<uint16_t>((lat - LAT_MIN) / (LAT_MAX - LAT_MIN) * 65535.0);
}

static inline uint8_t parse_speed_val(const char* s) {
    if (!s || !*s) return 0;
    char* end = nullptr;
    double v = strtod(s, &end);
    if (end != s && v > 0 && v <= 255) return static_cast<uint8_t>(v);
    if (strstr(s, "mph")) {
        v = strtod(s, &end);
        if (end != s) {
            v *= 1.609;
            if (v > 0 && v <= 255) return static_cast<uint8_t>(v);
        }
    }
    return 0;
}

class StreamHandler : public osmium::handler::Handler {
public:
    FILE* out;
    uint64_t way_count = 0;
    uint64_t point_count = 0;

    explicit StreamHandler(FILE* f) : out(f) {}

    void way(const osmium::Way& way) {
        const char* highway = way.tags().get_value_by_key("highway");
        if (!highway) return;

        const char* ms = way.tags().get_value_by_key("maxspeed");
        uint8_t speed = ms ? parse_speed_val(ms) : 0;
        if (speed == 0) return;

        way_count++;

        for (const auto& wn : way.nodes()) {
            if (!wn.location().valid()) continue;

            double lon = wn.location().lon();
            double lat = wn.location().lat();

            if (lon < LON_MIN || lon >= LON_MAX || lat < LAT_MIN || lat >= LAT_MAX) continue;

            RawPoint rp;
            rp.col = static_cast<int32_t>((lon - LON_MIN) / CELL);
            rp.row = static_cast<int32_t>((lat - LAT_MIN) / CELL);
            rp.lonq = quantize_lon(lon);
            rp.latq = quantize_lat(lat);
            rp.speed = speed;

            fwrite(&rp, sizeof(RawPoint), 1, out);
            point_count++;

            if (point_count % 5000000 == 0) {
                fprintf(stderr, "  %lu points (%lu ways)...\n", point_count, way_count);
            }
        }
    }
};

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

    const char* input = argv[1];
    const char* outpath = argv[2];

    fs::create_directories(fs::path(outpath).parent_path());

    FILE* out = fopen(outpath, "wb");
    if (!out) { fprintf(stderr, "Cannot open %s\n", outpath); return 1; }

    std::string idx_path = std::string(outpath) + ".idx";
    int idx_fd = ::open(idx_path.c_str(), O_CREAT | O_RDWR | O_TRUNC, 0644);
    if (idx_fd == -1) { fprintf(stderr, "Cannot create index\n"); fclose(out); return 1; }

    index_type index(idx_fd);
    location_handler_type location_handler(index);
    StreamHandler handler(out);

    fprintf(stderr, "Reading %s...\n", input);
    osmium::io::Reader reader(input);
    osmium::apply(reader, location_handler, handler);
    reader.close();

    fclose(out);
    ::close(idx_fd);
    unlink(idx_path.c_str());

    fprintf(stderr, "Ways with maxspeed: %lu\n", handler.way_count);
    fprintf(stderr, "Points written: %lu\n", handler.point_count);
    fprintf(stderr, "Output: %s\n", outpath);
    return 0;
}