#include "c5_radio.hpp"
#include "otm_tx_custom.h"
#include <vanetza_idf/its_g5_frame.hpp>
#include <esp_event.h>
#include <esp_log.h>
#include <esp_wifi.h>
#include <hal/modem_syscon_ll.h>
#include <esp_timer.h>
#include <freertos/FreeRTOS.h>
#include <freertos/queue.h>
#include <algorithm>
#include <atomic>
#include <cmath>
#include <cstring>
#include <iterator>
#include <limits>
#include <mutex>

extern "C" { //all of these arentt in the esp-idf public api
// phy_11p_set/phy_change_channel: undocumented esp_phy/lib/esp32c5/libphy.a entry points, not
// declared in any Espressif header. 
// The call sites and argument values below (phy_11p_set(1, 0), phy_change_channel(freq, 1, 0, 0))
// are copied from OpenTrafficMap's its-g5-receiver-firmware_txenabled, main/cmd_sniffer.c
// (https://codeberg.org/opentrafficmap/its-g5-receiver-firmware_txenabled, community reverse
// engineering, no stated license). 

void phy_11p_set(int enable, int arg2);
void phy_change_channel(int freq_mhz, int arg2_ignored, int arg3_ignored, int arg4);

// phy_get_cca/phy_set_cca: register 0x600a701c[7:0] holds the configured CCA energy
// detection threshold (defaults to 191 = 0xBF = -65 dBm in 8-bit two's complement).
// phy_get_cca() reads this configured threshold.
int phy_get_cca(void);
void phy_set_cca(int enable, int threshold);

// phy_get_cca_cnt/phy_set_cca_cnt: register 0x600a7c58 arms the 27-bit hardware CCA
// cycle counters (0x600a7c5c = total cycles, 0x600a7c60 = busy cycles; confirmed on
// hardware -- out[0] free-runs at ~40 MHz, out[1] stays near zero on a quiet channel).
// phy_get_cca_cnt returns bit 27 (busy/status bit) and writes both counters to out[2].
int phy_get_cca_cnt(std::int32_t out[2]);
void phy_set_cca_cnt(std::int32_t val, bool enable);
void phy_enable_cca(void);
void phy_disable_cca(void);
int phy_get_noise_floor(void);
}

namespace microbu {
namespace {
const char* TAG = "c5_radio";

}

class C5Radio::Impl {
public:
    /// @brief Raw received frame metadata and bytes from the promiscuous RX callback.
    struct Raw {
        std::uint16_t length;
        std::int8_t rssi;
        std::uint32_t timestamp;
        std::uint8_t bytes[2346]; // max 802.11 frame size
    };
    C5RadioConfig config;
    QueueHandle_t queue = nullptr;
    bool initialized = false, started = false, own_event_loop = false;
    std::uint16_t sequence = 0;
    std::atomic<std::uint32_t> dropped {0};
    static Impl* active;
    static std::mutex callback_mutex;

    explicit Impl(C5RadioConfig c) : config(c) {}

    static void receive(void* buffer, wifi_promiscuous_pkt_type_t type) {
        if (!buffer || type != WIFI_PKT_DATA) return;
        const auto* packet = static_cast<const wifi_promiscuous_pkt_t*>(buffer);
        if (packet->rx_ctrl.rx_state != 0) return;
        const auto length = packet->rx_ctrl.sig_len;
        std::lock_guard<std::mutex> lock(callback_mutex);
        if (!active || !active->queue) return;
        // MicrOBU: a frame too short to hold any GN packet is not ITS traffic, so it is ignored
        // rather than counted; dropped_frames() then means what the phone shows it as: lost frames.
        if (length < 38) return;
        if (length > sizeof(Raw::bytes)) { ++active->dropped; return; }
        // MicrOBU: static, not a local. Raw is ~2.4 KB and this runs on the Wi-Fi driver's own task,
        // several frames deep, on a stack of roughly 3.5 KB: the previous firmware hit exactly this
        // with an 800-byte buffer (obu-firmware commit 04b0076). Safe as a static because only that
        // one task calls this, and xQueueSend copies it out before the next call.
        static Raw raw;
        raw.length = length;
        raw.rssi = packet->rx_ctrl.rssi;
        raw.timestamp = packet->rx_ctrl.timestamp;
        std::memcpy(raw.bytes, packet->payload, length);
        if (xQueueSend(active->queue, &raw, 0) != pdTRUE) ++active->dropped;
    }
};

C5Radio::Impl* C5Radio::Impl::active = nullptr;
std::mutex C5Radio::Impl::callback_mutex;

C5Radio::C5Radio(C5RadioConfig c) : impl_(std::make_unique<Impl>(c)) {}
C5Radio::~C5Radio() { stop(); }

esp_err_t C5Radio::start() {
    auto& p = *impl_;
    const auto& c = p.config;
    if (p.initialized) return ESP_ERR_INVALID_STATE;
    // Reject channel/power/queue config outside the supported ITS-G5 range
    if (c.channel_number < 172 || c.channel_number > 184 || c.channel_number % 2 ||
        !std::isfinite(c.transmit_power_dbm) || c.transmit_power_dbm < 2 || c.transmit_power_dbm > 23 ||
        std::floor(c.transmit_power_dbm * 4) != c.transmit_power_dbm * 4 ||
        c.receive_queue_length == 0 || c.receive_queue_length > 32) {
        return ESP_ERR_INVALID_ARG;
    }

    {
        std::lock_guard<std::mutex> lock(Impl::callback_mutex);
        if (Impl::active) return ESP_ERR_INVALID_STATE;
        p.queue = xQueueCreate(c.receive_queue_length, sizeof(Impl::Raw));
        if (!p.queue) return ESP_ERR_NO_MEM;
        Impl::active = &p;
    }

    auto result = esp_event_loop_create_default();
    p.own_event_loop = (result == ESP_OK);
    if (result != ESP_OK && result != ESP_ERR_INVALID_STATE) {
        stop();
        return result;
    }

    // Establish modem FE clock for 802.11p OFDM
    modem_syscon_ll_enable_fe_40m_clock(&MODEM_SYSCON, true);

    wifi_init_config_t wifi = WIFI_INIT_CONFIG_DEFAULT();
    wifi.nvs_enable = 0;
    result = esp_wifi_init(&wifi);
    if (result != ESP_OK) {
        stop();
        return result;
    }
    p.initialized = true;

    auto attempt = [&](esp_err_t r) { if (result == ESP_OK) result = r; };
    attempt(esp_wifi_set_storage(WIFI_STORAGE_RAM));
    attempt(esp_wifi_set_mode(WIFI_MODE_STA));
    if (result == ESP_OK) {
        result = esp_wifi_start();
        p.started = (result == ESP_OK);
    }
    if (result != ESP_OK) {
        stop();
        return result;
    }

    // MicrOBU, from the previous firmware: under the default WIFI_COUNTRY_POLICY_AUTO the driver's
    // 5 GHz table does not authorise transmission on the ITS band, which there left RX working and
    // TX silent. The colleague's board transmits without this (esp_wifi_80211_tx_custom goes around
    // that gate), so it is belt and braces here: manual policy, every 5 GHz channel enabled. Not
    // fatal if refused.
    wifi_country_t country = {};
    country.cc[0] = 'U'; country.cc[1] = 'S';
    country.schan = 1;
    country.nchan = 11;
    country.policy = WIFI_COUNTRY_POLICY_MANUAL;
    country.wifi_5g_channel_mask = 0x1FFFFFFE;
    if (const auto e = esp_wifi_set_country(&country); e != ESP_OK)
        ESP_LOGW(TAG, "esp_wifi_set_country(MANUAL) failed: %s (continuing)", esp_err_to_name(e));
    attempt(esp_wifi_set_band_mode(WIFI_BAND_MODE_5G_ONLY));
    attempt(esp_wifi_set_ps(WIFI_PS_NONE));
    attempt(esp_wifi_set_max_tx_power(static_cast<std::int8_t>(c.transmit_power_dbm * 4)));

    wifi_promiscuous_filter_t filter {};
    filter.filter_mask = WIFI_PROMIS_FILTER_MASK_DATA;
    attempt(esp_wifi_set_promiscuous_filter(&filter));
    attempt(esp_wifi_set_promiscuous_rx_cb(Impl::receive));
    attempt(esp_wifi_set_promiscuous(true));
    if (result != ESP_OK) {
        stop();
        return result;
    }

    // 10 MHz channel bandwidth (ITS-G5 / 802.11p)
    phy_11p_set(1, 0);
    phy_change_channel(5000 + 5 * c.channel_number, 1, 0, 0); // = 5900 MHz

    // Enable and arm hardware CCA counters (40 MHz baseband clock timebase) for DCC
    phy_enable_cca();
    phy_set_cca_cnt(0x07FFFFFF, true);

    // MicrOBU, from the previous firmware: the power request is a ceiling, not a promise. The driver
    // clamps it to its calibrated table, and 5900 MHz is above the chip's rated range, so log what
    // the driver admits to rather than what was asked for.
    std::int8_t power_q = 0;
    if (esp_wifi_get_max_tx_power(&power_q) == ESP_OK) {
        ESP_LOGI(TAG, "tx power: %d quarter-dBm = %d.%02d dBm (%.2f requested)",
                 power_q, power_q / 4, (power_q % 4) * 25, c.transmit_power_dbm);
    }

    ESP_LOGI(TAG, "ITS-G5 802.11p radio started on channel %u (5900 MHz), %s",
             unsigned(c.channel_number), c.laboratory_transmission ? "TX/RX" : "RX only");
    return ESP_OK;
}

void C5Radio::stop() {
    if (!impl_) return;
    auto& p = *impl_;
    if (p.started) esp_wifi_set_promiscuous(false);
    {
        std::lock_guard<std::mutex> lock(Impl::callback_mutex);
        if (Impl::active == &p) Impl::active = nullptr;
        if (p.queue) {
            vQueueDelete(p.queue);
            p.queue = nullptr;
        }
    }
    if (p.started) esp_wifi_stop();
    if (p.initialized) esp_wifi_deinit();
    if (p.own_event_loop) esp_event_loop_delete_default();
    p.started = p.initialized = p.own_event_loop = false;
}

vanetza_idf::Result C5Radio::request(vanetza_idf::AlDataRequest request) {
    auto& p = *impl_;
    if (!p.started) return vanetza_idf::Result::rejected;
    if (!p.config.laboratory_transmission) return vanetza_idf::Result::unsupported;

    if (request.bandwidth_mhz != 10 || request.channel_number != p.config.channel_number ||
        request.transceiver_id != 0 || request.transceiver_mode || request.datastream_id ||
        request.transmit_power_dbm != p.config.transmit_power_dbm) {
        return vanetza_idf::Result::unsupported;
    }

    constexpr wifi_phy_rate_t rates[] = {
        WIFI_PHY_RATE_6M, WIFI_PHY_RATE_9M, WIFI_PHY_RATE_12M,
        WIFI_PHY_RATE_18M, WIFI_PHY_RATE_24M, WIFI_PHY_RATE_36M,
        WIFI_PHY_RATE_48M, WIFI_PHY_RATE_54M
    };
    const auto index = static_cast<unsigned>(request.mcs);
    if (index >= std::size(rates)) return vanetza_idf::Result::invalid_argument;

    vanetza::ByteBuffer bytes;
    const auto encoded = vanetza_idf::its_g5::encode_frame(request, p.sequence, bytes);
    if (encoded != vanetza_idf::Result::accepted) {
        ESP_LOGE(TAG, "encode_frame failed: %d", int(encoded));
        return encoded;
    }

    p.sequence = (p.sequence + 1) & 4095;
    wifi_tx_rate_config_t rate {};
    rate.phymode = WIFI_PHY_MODE_11A;
    rate.rate = rates[index];

    // Transmit frame via 802.11p driver
    const auto result = esp_wifi_80211_tx_custom(
        WIFI_IF_STA, bytes.data(), bytes.size(), false,
        &rate, WIFI_BAND_5G, WIFI_BW20);

    if (result != ESP_OK) {
        ESP_LOGW(TAG, "esp_wifi_80211_tx_custom failed: %s (0x%x)", esp_err_to_name(result), result);
    }

    return result == ESP_OK ? vanetza_idf::Result::accepted :
           result == ESP_ERR_NO_MEM ? vanetza_idf::Result::resource_limit :
           vanetza_idf::Result::rejected;
}

void C5Radio::poll(const Receive& receive, const Capture& capture) {
    auto& p = *impl_;
    if (!p.queue) return;
    Impl::Raw raw {};
    for (unsigned i = 0; i < p.config.receive_queue_length && xQueueReceive(p.queue, &raw, 0) == pdTRUE; ++i) {
        if (capture) {
            capture(vanetza::ByteBuffer(raw.bytes, raw.bytes + raw.length), raw.rssi, raw.timestamp);
        }
        vanetza_idf::AlDataIndication ind;
        if (vanetza_idf::its_g5::decode_frame(raw.bytes, raw.length, true, ind) != vanetza_idf::Result::accepted) {
            continue;
        }
        ind.channel_number = p.config.channel_number;
        ind.received_power_dbm = raw.rssi;
        if (receive) receive(std::move(ind));
    }
}

std::uint32_t C5Radio::dropped_frames() const {
    return impl_->dropped.load();
}

#if CONFIG_MICROBU_TEST_CHANNEL
esp_err_t C5Radio::transmit_burst(std::uint16_t channel, double power_dbm, unsigned mcs,
                                  unsigned count, unsigned interval_ms, std::size_t payload_len) {
    auto& p = *impl_;
    if (!p.started) return ESP_ERR_INVALID_STATE;
    if (channel < 172 || channel > 184 || channel % 2 != 0) return ESP_ERR_INVALID_ARG;
    if (power_dbm < 2.0 || power_dbm > 20.0) return ESP_ERR_INVALID_ARG;
    if (mcs > 7) return ESP_ERR_INVALID_ARG;
    if (count == 0) return ESP_OK;

    constexpr wifi_phy_rate_t rates[] = {
        WIFI_PHY_RATE_6M, WIFI_PHY_RATE_9M, WIFI_PHY_RATE_12M,
        WIFI_PHY_RATE_18M, WIFI_PHY_RATE_24M, WIFI_PHY_RATE_36M,
        WIFI_PHY_RATE_48M, WIFI_PHY_RATE_54M
    };

    // Dynamically retune channel or adjust TX power if different from running config
    if (channel != p.config.channel_number) {
        phy_11p_set(1, 0);
        phy_change_channel(5000 + 5 * channel, 1, 0, 0);
        p.config.channel_number = channel;
    }
    const auto power_quarter_db = static_cast<std::int8_t>(std::round(power_dbm * 4.0)); // esp_wifi power is in 0.25 dBm units
    esp_wifi_set_max_tx_power(power_quarter_db);
    p.config.transmit_power_dbm = power_dbm;

    wifi_tx_rate_config_t rate {};
    rate.phymode = WIFI_PHY_MODE_11A;
    rate.rate = rates[mcs];

    // Assemble a standard IEEE 802.11 QoS data / LLC frame (EtherType 0x8947 GeoNetworking)
    // Header: Frame Control (0x0088 QoS Data), Duration (0x0000), Addr1 (Broadcast FF..FF),
    // Addr2 (Source 02:00:00:00:00:01), Addr3 (BSSID FF..FF), Sequence, QoS Control (0x0000),
    // LLC/SNAP header (AA AA 03 00 00 00 89 47).
    std::vector<std::uint8_t> frame;
    const std::size_t actual_payload = std::clamp<std::size_t>(payload_len, 32, 1400);
    frame.reserve(34 + actual_payload);

    // MAC Header (26 bytes with QoS)
    frame.push_back(0x88); frame.push_back(0x00); // Frame Control: QoS Data
    frame.push_back(0x00); frame.push_back(0x00); // Duration
    for (int i = 0; i < 6; ++i) frame.push_back(0xFF); // RA / Destination: Broadcast
    frame.push_back(0x02); frame.push_back(0x00); frame.push_back(0x00);
    frame.push_back(0x00); frame.push_back(0x00); frame.push_back(0x01); // TA / Source
    for (int i = 0; i < 6; ++i) frame.push_back(0xFF); // BSSID: Broadcast
    frame.push_back(0x00); frame.push_back(0x00); // Sequence (updated per frame)
    frame.push_back(0x00); frame.push_back(0x00); // QoS Control

    // LLC/SNAP header (8 bytes)
    frame.push_back(0xAA); frame.push_back(0xAA); frame.push_back(0x03);
    frame.push_back(0x00); frame.push_back(0x00); frame.push_back(0x00);
    frame.push_back(0x89); frame.push_back(0x47); // EtherType 0x8947 (GeoNetworking)

    // Test payload with identifiable sequence numbers
    const std::size_t header_len = frame.size();
    frame.resize(header_len + actual_payload, 0x5A);

    esp_err_t last_err = ESP_OK;
    for (unsigned i = 0; i < count; ++i) {
        p.sequence = (p.sequence + 1) & 4095;
        frame[22] = static_cast<std::uint8_t>((p.sequence << 4) & 0xF0);
        frame[23] = static_cast<std::uint8_t>((p.sequence >> 4) & 0xFF);

        // Put burst counter inside payload
        frame[header_len + 0] = static_cast<std::uint8_t>(i & 0xFF);
        frame[header_len + 1] = static_cast<std::uint8_t>((i >> 8) & 0xFF);

        esp_err_t err = esp_wifi_80211_tx_custom(
            WIFI_IF_STA, frame.data(), frame.size(), false,
            &rate, WIFI_BAND_5G, WIFI_BW20);

        if (err == ESP_ERR_NO_MEM) {
            // Buffer briefly full: yield task to allow DMA descriptors to clear
            vTaskDelay(pdMS_TO_TICKS(2));
            err = esp_wifi_80211_tx_custom(
                WIFI_IF_STA, frame.data(), frame.size(), false,
                &rate, WIFI_BAND_5G, WIFI_BW20);
        }

        if (err != ESP_OK) {
            last_err = err;
        }

        const auto delay_ms = std::max<unsigned>(interval_ms, 2);
        if (i + 1 < count) {
            vTaskDelay(pdMS_TO_TICKS(delay_ms));
        }
    }

    return last_err;
}

CcaSampleResult C5Radio::sample_cca(unsigned duration_ms) {
    CcaSampleResult result;
    auto& p = *impl_;
    if (!p.started) return result;

    result.noise_floor_dbm = phy_get_noise_floor();

    // Enable CCA hardware and arm the 27-bit cycle counters with full window (0x07FFFFFF).
    phy_enable_cca();
    phy_set_cca_cnt(0x07FFFFFF, true);

    std::int32_t cca_cnt_before[2] = {};
    phy_get_cca_cnt(cca_cnt_before);

    const auto t_start = esp_timer_get_time();
    const auto t_deadline = t_start + static_cast<std::int64_t>(duration_ms) * 1000;
    std::int64_t last_t = t_start;
    result.min_delta_us = std::numeric_limits<std::uint32_t>::max();

    while (esp_timer_get_time() < t_deadline) {
        std::int32_t cur_cnt[2] = {};
        const auto status = phy_get_cca_cnt(cur_cnt);
        const auto cca_threshold = phy_get_cca();
        const auto now = esp_timer_get_time();

        if (result.samples == 0) {
            result.first_cca = cca_threshold;
        } else {
            const auto delta = static_cast<std::uint32_t>(now - last_t);
            result.min_delta_us = std::min(result.min_delta_us, delta);
            result.max_delta_us = std::max(result.max_delta_us, delta);
        }
        last_t = now;
        result.last_cca = cca_threshold;
        if (status) ++result.busy_count;
        ++result.samples;

        // Cooperative yielding: prevent starving IDLE task, esp_timer, and bb_wdt
        // on the single-core C5 during multi-millisecond polling windows.
        if ((result.samples & 0x3F) == 0) {
            taskYIELD();
        }
    }

    result.duration_us = static_cast<std::uint32_t>(esp_timer_get_time() - t_start);
    std::int32_t cca_cnt_after[2] = {};
    result.cca_status = phy_get_cca_cnt(cca_cnt_after);

    // Both words are 27-bit hardware counters (mask 0x07FFFFFF).
    constexpr std::int32_t mask27 = 0x07FFFFFF;
    auto delta27 = [](std::int32_t after, std::int32_t before) -> std::int32_t {
        std::int32_t diff = (after & mask27) - (before & mask27);
        if (diff < 0) diff += (mask27 + 1);
        return diff;
    };
    result.cca_total_cycles_delta = delta27(cca_cnt_after[0], cca_cnt_before[0]);
    result.cca_busy_cycles_delta = delta27(cca_cnt_after[1], cca_cnt_before[1]);

    if (result.samples < 2) result.min_delta_us = 0;
    return result;
}
#endif // CONFIG_MICROBU_TEST_CHANNEL

CcaCounters C5Radio::read_cca_counters() const {
    CcaCounters c;
    std::int32_t out[2] = {};
    phy_get_cca_cnt(out);
    constexpr std::int32_t mask27 = 0x07FFFFFF; //  remomve the busy/status bit (bit 27) from the 27-bit hardware counters
    c.total_cycles = static_cast<std::uint32_t>(out[0] & mask27);
    c.busy_cycles = static_cast<std::uint32_t>(out[1] & mask27);
    return c;
}

double C5Radio::calculate_cbr(const CcaCounters& current, const CcaCounters& previous) {
    constexpr std::uint32_t counter_range = 1u << 27; 
    auto delta27 = [](std::uint32_t after, std::uint32_t before) -> std::uint32_t {
        if (after >= before) {
            return after - before;
        }
        // Counter wrapped from 2^27 - 1 back to zero.
        return after + counter_range - before;
    };

    const std::uint32_t dt = delta27(current.total_cycles, previous.total_cycles);
    const std::uint32_t db = delta27(current.busy_cycles, previous.busy_cycles);

    if (dt == 0) return 0.0;
    return static_cast<double>(db) / static_cast<double>(dt);
}

} // namespace microbu