#include <unistd.h>
#include <cstdlib>
#include <iostream>
#include <pcap.h>
#include <cinttypes>
#include <sstream>
#include <iomanip>
#include "../third-party/radiotap-library/radiotap.h"
#include "wifi-triangulator/core.h"

using namespace std;
using namespace wifi_triangulator::pb;

namespace wifi_triangulator {

std::string app_name;

struct ieee802_11_header {
    uint8_t frame_1;
    uint8_t frame_2;
    uint16_t duration_id;

    eth_mac addr1, addr2, addr3;
    uint16_t seq_ctl;
    eth_mac addr4;

    bool is_to_ds() const {
        return (frame_2 & 0x80) > 0;
    }

    bool is_from_ds() const {
        return (frame_2 & 0x40) > 0;
    }
} __attribute__ ((packed));

struct radio_tap_it {
    radio_tap_it(const uint8_t *data) : offset(0), data(data) {
    }

    const uint8_t *data;
    int offset;

    uint8_t read8u() {
        uint8_t value = data[offset];
        offset += 1;
        return value;
    }

    int8_t read8s() {
        int8_t value = data[offset];
        offset += 1;
        return value;
    }

    uint16_t read16u() {
        uint16_t value = *reinterpret_cast<const uint16_t *>(&data[offset]);
        offset += 2;
        offset = (offset + (2 - 1)) & -2;
        return value;
    }

    uint32_t read32u() {
        uint32_t value = *reinterpret_cast<const uint32_t *>(&data[offset]);
        offset += 4;
        offset = (offset + (4 - 1)) & -4;
        return value;
    }

    uint64_t read64u() {
        uint64_t value = *reinterpret_cast<const uint64_t *>(&data[offset]);
        offset += 8;
        offset = (offset + (8 - 1)) & -8;
        return value;
    }
};

struct capture_context {
    std::function<void(const data &)> data_consumer;
};

void got_packet(u_char *args, const struct pcap_pkthdr *header, const u_char *packet) {
//    fprintf(stdout, "%lu.%lu, caplen=%d, len=%d\n", header->ts.tv_sec, header->ts.tv_usec, header->caplen, header->len);

    auto *ctx = reinterpret_cast<capture_context *>(args);
    auto *rtaphdr = reinterpret_cast<const struct ieee80211_radiotap_header *>(packet);

    if (rtaphdr->it_version != 0) {
        return;
    }

    int present_count = 0;
    {
        auto *present_ptr = &rtaphdr->it_present;
        do {
            present_count++;
        } while ((*present_ptr++) & 1 << IEEE80211_RADIOTAP_EXT);
    }

    radio_tap_it it(packet + 4 + 4 * present_count);

    auto *present_ptr = &rtaphdr->it_present;
    bool is_radiotap = true;
    uint16_t mhz = 0;
    int rssi = 0;
    do {
        uint32_t present = *present_ptr;
        bool next_is_radiotap = (present & 1 << IEEE80211_RADIOTAP_RADIOTAP_NAMESPACE) != 0;
        // printf("present: %08x\n", present);

        present &= ~(1 << IEEE80211_RADIOTAP_RADIOTAP_NAMESPACE);
        present &= ~(1 << IEEE80211_RADIOTAP_VENDOR_NAMESPACE);
        present &= ~(1 << IEEE80211_RADIOTAP_EXT);

        if (is_radiotap) {
            char buf[101];

            if (present & 1 << IEEE80211_RADIOTAP_TSFT) {
                auto tsft = it.read64u();
                snprintf(buf, 100, "IEEE80211_RADIOTAP_TSFT: %" PRIu64 "\n", tsft);
                present &= ~(1 << IEEE80211_RADIOTAP_TSFT);
            }
            if (present & 1 << IEEE80211_RADIOTAP_FLAGS) {
                uint8_t flags = it.read8u();
                snprintf(buf, 100, "IEEE80211_RADIOTAP_FLAGS: %02x\n", flags);
                present &= ~(1 << IEEE80211_RADIOTAP_FLAGS);
            }
            if (present & 1 << IEEE80211_RADIOTAP_RATE) {
                uint8_t rate = it.read8u();

                double r = rate / 2;
                snprintf(buf, 100, "IEEE80211_RADIOTAP_RATE: %.1lf\n", r);
                present &= ~(1 << IEEE80211_RADIOTAP_RATE);
            }
            if (present & 1 << IEEE80211_RADIOTAP_CHANNEL) {
                mhz = it.read16u();
                uint16_t bitmap = it.read16u();

                snprintf(buf, 100, "IEEE80211_RADIOTAP_CHANNEL: %d MHz, flags=%04x\n", mhz, bitmap);
                present &= ~(1 << IEEE80211_RADIOTAP_CHANNEL);
            }
            if (present & 1 << IEEE80211_RADIOTAP_FHSS) {
                uint8_t hop_set = it.read8u();
                uint8_t pattern = it.read8u();

                snprintf(buf, 100, "IEEE80211_RADIOTAP_FHSS: hop_set=%02x, pattern=%02x\n", hop_set, pattern);
                present &= ~(1 << IEEE80211_RADIOTAP_FHSS);
            }
            if (present & 1 << IEEE80211_RADIOTAP_DBM_ANTSIGNAL) {
                int8_t antsignal_dbm = it.read8s();

                rssi = antsignal_dbm;
                snprintf(buf, 100, "IEEE80211_RADIOTAP_DBM_ANTSIGNAL: dbm=%d\n", antsignal_dbm);
                present &= ~(1 << IEEE80211_RADIOTAP_DBM_ANTSIGNAL);
            }
            if (present & 1 << IEEE80211_RADIOTAP_DBM_ANTNOISE) {
                int8_t antnoise_dbm = it.read8s();

                snprintf(buf, 100, "IEEE80211_RADIOTAP_DBM_ANTNOISE: dbm=%d\n", antnoise_dbm);
                present &= ~(1 << IEEE80211_RADIOTAP_DBM_ANTNOISE);
            }
            if (present & 1 << IEEE80211_RADIOTAP_LOCK_QUALITY) {
                uint16_t lock_quality = it.read16u();

                snprintf(buf, 100, "IEEE80211_RADIOTAP_LOCK_QUALITY: lock_quality=%d\n", lock_quality);
                present &= ~(1 << IEEE80211_RADIOTAP_LOCK_QUALITY);
            }
            if (present & 1 << IEEE80211_RADIOTAP_TX_ATTENUATION) {
                uint16_t tx_attenuation = it.read16u();

                snprintf(buf, 100, "IEEE80211_RADIOTAP_TX_ATTENUATION: tx_attenuation=%d\n", tx_attenuation);
                present &= ~(1 << IEEE80211_RADIOTAP_TX_ATTENUATION);
            }
            if (present & 1 << IEEE80211_RADIOTAP_DB_TX_ATTENUATION) {
                uint16_t db_tx_attenuation = it.read16u();

                snprintf(buf, 100, "IEEE80211_RADIOTAP_DB_TX_ATTENUATION: db_tx_attenuation=%d\n", db_tx_attenuation);
                present &= ~(1 << IEEE80211_RADIOTAP_DB_TX_ATTENUATION);
            }
            if (present & 1 << IEEE80211_RADIOTAP_DBM_TX_POWER) {
                int8_t tx_power_dbm = it.read8s();

                snprintf(buf, 100, "IEEE80211_RADIOTAP_DBM_TX_POWER: tx_power_dbm=%d\n", tx_power_dbm);
                present &= ~(1 << IEEE80211_RADIOTAP_DBM_TX_POWER);
            }
            if (present & 1 << IEEE80211_RADIOTAP_ANTENNA) {
                uint8_t antenna = it.read8u();

//                snprintf(buf, 100, "IEEE80211_RADIOTAP_ANTENNA: antenna=%d\n", antenna);
                present &= ~(1 << IEEE80211_RADIOTAP_ANTENNA);
            }
            if (present & 1 << IEEE80211_RADIOTAP_DB_ANTSIGNAL) {
                uint8_t antenna_signal_db = it.read8u();

                snprintf(buf, 100, "IEEE80211_RADIOTAP_DB_ANTSIGNAL: antenna_signal_db=%d\n", antenna_signal_db);
                present &= ~(1 << IEEE80211_RADIOTAP_DB_ANTSIGNAL);
            }
            if (present & 1 << IEEE80211_RADIOTAP_DB_ANTNOISE) {
                uint8_t antenna_noise_db = it.read8u();

                snprintf(buf, 100, "IEEE80211_RADIOTAP_DB_ANTNOISE: antenna_noise_db=%d\n", antenna_noise_db);
                present &= ~(1 << IEEE80211_RADIOTAP_DB_ANTNOISE);
            }
            if (present & 1 << IEEE80211_RADIOTAP_RX_FLAGS) {
                uint16_t rx_flags = it.read16u();

//                snprintf(buf, 100, "IEEE80211_RADIOTAP_RX_FLAGS: rx_flags=%d\n", rx_flags);
                present &= ~(1 << IEEE80211_RADIOTAP_RX_FLAGS);
            }
            if (present & 1 << IEEE80211_RADIOTAP_TX_FLAGS) {
                uint16_t tx_flags = it.read16u();

                snprintf(buf, 100, "IEEE80211_RADIOTAP_TX_FLAGS: tx_flags=%d\n", tx_flags);
                present &= ~(1 << IEEE80211_RADIOTAP_TX_FLAGS);
            }

            if (present > 0) {
                snprintf(buf, 100, "Unknown fields remain: %08x\n", present);
                break;
            }

//            fprintf(stderr, buf);
//            fflush(stderr);
            /*
            int bit = 0;
            uint32_t mask = 1;
            for (int i = 0; i < 29; i++) {
                if (present & mask) {
                    printf("bit %d\n", bit);
                }

                bit++;
                mask <<= 1;
            }
            */
        }
        is_radiotap = next_is_radiotap;
    } while ((*present_ptr++) & 1 << IEEE80211_RADIOTAP_EXT);

    const struct ieee802_11_header *ieee802_11_header = reinterpret_cast<const struct ieee802_11_header *>(&packet[rtaphdr->it_len]);

    uint8_t type = ieee802_11_header->frame_1;

    frame_type t;
    switch (type) {
        case 0x40:
            t = frame_type::probe_request;
            break;
        case 0x80:
            t = frame_type::beacon;
            break;
        case 0x50:
            t = frame_type::probe_response;
            break;
        case 0x48:
            t = frame_type::null;
            break;
        case 0xd4:
            t = frame_type::ack;
            break;
        case 0x08:
            t = frame_type::data;
            break;
        case 0xc4:
            t = frame_type::cts;
            break;
        case 0xb4:
            t = frame_type::rts;
            break;
        case 0x1e:
            t = frame_type::cf_end;
            break;
        case 0x1f:
            t = frame_type::cf_end_cf_ack;
            break;
        case 0x1a:
            t = frame_type::ps_poll;
            break;
        default:
            t = frame_type::unknown;
    }

//    printf("ieee802_11_header->frame_ctl=%02x, %s\n", type, type_str ? type_str : "???");

    eth_mac src, dst;

    if (!ieee802_11_header->is_from_ds()) {
        if (!ieee802_11_header->is_to_ds()) {
            dst = ieee802_11_header->addr1;
            src = ieee802_11_header->addr2;
        } else {
            dst = ieee802_11_header->addr1;
            src = ieee802_11_header->addr3;
        }
    } else {
        if (!ieee802_11_header->is_to_ds()) {
            dst = ieee802_11_header->addr3;
            src = ieee802_11_header->addr2;
        } else {
            dst = ieee802_11_header->addr3;
            src = ieee802_11_header->addr4;
        }
    }

    ctx->data_consumer(data{
            t,
            header->ts.tv_sec, header->ts.tv_usec,
            rssi, mhz,
            src, dst});
}

int launch_capture(string dev, std::function<void(const class data &)> data_consumer) {
    char errbuf[1000];

    pcap_t *handle;

    handle = pcap_open_live(dev.c_str(), BUFSIZ, 1, 1000, errbuf);
//    handle = pcap_open_offline(dev, errbuf);

    if (!handle) {
        fprintf(stderr, "Could not open %s: %s\n", dev.c_str(), errbuf);
        return EXIT_FAILURE;
    }

    int i = pcap_datalink(handle);
    if (i != DLT_IEEE802_11_RADIO) {
        fprintf(stderr, "Device %s doesn't provide IEEE 802.11 radio headers\n", dev.c_str());
        return EXIT_FAILURE;
    }

/* Compile and apply the filter */
//    if (pcap_compile(handle, &fp, filter_exp, 0, net) == -1) {
//        fprintf(stderr, "Couldn't parse filter %s: %s\n", filter_exp, pcap_geterr(handle));
//        return EXIT_FAILURE;
//    }
//    if (pcap_setfilter(handle, &fp) == -1) {
//        fprintf(stderr, "Couldn't install filter %s: %s\n", filter_exp, pcap_geterr(handle));
//        return EXIT_FAILURE;
//    }

    capture_context ctx{data_consumer};
    while (true) {
        auto ret = pcap_loop(handle, 1000, got_packet, reinterpret_cast<u_char *>(&ctx));

        if (ret == -1) {
            printf("pcap failed: %s\n", pcap_geterr(handle));
            break;
        } else if (ret == -2) {
            // someone called break loop
            break;
        }
    }

    pcap_close(handle);

    return EXIT_SUCCESS;
}

int launch_reader(std::function<bool(uint16_t size, const uint8_t *data)> on_buffer) {
    uint8_t buffer[1 << 16];
    while (true) {
        uint16_t size;

        if (read_stdin(reinterpret_cast<uint8_t *>(&size), 2) < 0) {
            break;
        }

        ssize_t got;
        if ((got = read_stdin(buffer, size)) != size) {
            cout << "short read, got=" << got << endl;
            return EXIT_FAILURE;
        }

        bool ok = on_buffer(size, buffer);

        if (!ok) {
            break;
        }
    }

    return EXIT_SUCCESS;
}

int launch_reader_envelope(std::function<bool(const pb::envelope &)> on_envelope) {
    pb::envelope envelope;

    return launch_reader([&](uint16_t size, const uint8_t *data) {
        bool ok = envelope.ParseFromArray(data, size);
        return ok && on_envelope(envelope);
    });
}

ssize_t read_stdin(uint8_t *data, const size_t count) {
    auto left = count;
    while (left > 0) {
        ssize_t n_read = read(STDIN_FILENO, data, left);

        if (n_read < 0) {
            return n_read;
        }

        left -= n_read;
    }

    return count;
}

//void write_stdout(const bstring &s) {
//    write_stdout(s.c_str(), s.length());
//}

void write_stdout(const std::string &s) {
    static_assert(sizeof(char) == sizeof(uint8_t), "bad sizes");
    write_stdout(reinterpret_cast<const uint8_t *>(s.c_str()), s.length());
}

void write_stdout(const uint8_t *data, size_t count) {
    size_t left = count;
    const uint8_t *buf = data;
    while (left > 0) {
        auto written = write(STDOUT_FILENO, data, left);
        buf += written;
        left -= written;
    }
    syncfs(STDOUT_FILENO);
}

void write_envelope(pb::letter_type lt, pb::envelope &envelope) {
    std::string str;
    str.reserve(64 * 1024);

    envelope.set_type(lt);

    envelope.SerializeToString(&str);
    uint16_t len = static_cast<uint16_t>(str.length());
    write_stdout(reinterpret_cast<uint8_t *>(&len), 2);
    write_stdout(str);
}

} // namespace wifi_triangulator