#include "../librevna_minimal_driver.hpp"

#include <array>
#include <complex>
#include <cstddef>
#include <stdexcept>
#include <vector>

#include "librevna_protocol_common.hpp"

namespace radar::drivers {
namespace detail = radar::drivers::librevna::detail;

auto LibreVnaMinimalDriver::encode_frame(
    std::uint8_t packet_type,
    std::span<const std::uint8_t> payload
) -> std::vector<std::uint8_t> {
    const auto length = payload.size() + detail::kFrameOverheadBytes;
    if (length > 0xFFFFU) {
        throw std::runtime_error("Protocol frame is too large");
    }

    std::vector<std::uint8_t> frame(length, 0);
    frame[0] = detail::kFrameHeader;
    detail::write_u16_le(frame, 1, static_cast<std::uint16_t>(length));
    frame[3] = packet_type;
    std::copy(payload.begin(), payload.end(), frame.begin() + 4);

    std::uint32_t crc = 0;
    if (packet_type != detail::kPacketVnaDatapoint) {
        crc = crc32(std::span<const std::uint8_t>(frame.data(), frame.size() - 4U));
    }
    detail::write_u32_le(frame, frame.size() - 4U, crc);

    return frame;
}

auto LibreVnaMinimalDriver::decode_vna_datapoint_traces(
    std::span<const std::uint8_t> payload,
    DecodedVnaDatapoint& datapoint_out
) -> bool {
    // VNADatapoint payload layout:
    // [0..7]=freq_or_time, [8..9]=cdbm, [10..11]=point_number,
    // followed by N tuples of {real:f32, imag:f32, flags:u8}.
    if (payload.size() < 12U) {
        return false;
    }

    const auto values_block = payload.size() - 12U;
    if ((values_block % 9U) != 0U) {
        return false;
    }

    const auto num_values = values_block / 9U;
    if (num_values == 0U) {
        return false;
    }

    datapoint_out = DecodedVnaDatapoint{};
    datapoint_out.point_number = detail::read_u16_le(payload, 10);
    datapoint_out.frequency_hz = static_cast<float>(detail::read_u64_le(payload, 0));

    const auto real_offset = 12U;
    const auto imag_offset = real_offset + (4U * num_values);
    const auto flags_offset = imag_offset + (4U * num_values);

    std::array<std::complex<float>, 8> ref_by_stage{};
    std::array<std::complex<float>, 8> s11_measured_by_stage{};
    std::array<std::complex<float>, 8> s21_measured_by_stage{};
    std::array<bool, 8> has_ref{};
    std::array<bool, 8> has_s11_measured{};
    std::array<bool, 8> has_s21_measured{};
    constexpr float kReferenceMagnitudeSquaredEpsilon = 1e-12F;

    for (std::size_t index = 0; index < num_values; ++index) {
        const auto flags = payload[flags_offset + index];
        const auto stage = static_cast<std::size_t>(flags >> 5U);
        if (stage >= ref_by_stage.size()) {
            continue;
        }

        const auto value = std::complex<float>(
            detail::read_f32_le(payload, real_offset + (4U * index)),
            detail::read_f32_le(payload, imag_offset + (4U * index))
        );

        const bool is_reference = (flags & detail::kReferenceFlagMask) != 0U;
        if ((flags & detail::kPort1Mask) != 0U && is_reference) {
            ref_by_stage[stage] = value;
            has_ref[stage] = true;
        }
        if ((flags & detail::kPort1Mask) != 0U && !is_reference) {
            s11_measured_by_stage[stage] = value;
            has_s11_measured[stage] = true;
        }
        if ((flags & detail::kPort2Mask) != 0U && !is_reference) {
            s21_measured_by_stage[stage] = value;
            has_s21_measured[stage] = true;
        }
    }

    // For the forward sweep used in this project, port1 reference is the
    // incident signal, port1 measured is reflection (S11 numerator), and
    // port2 measured is transmission (S21 numerator).
    for (std::size_t stage = 0; stage < ref_by_stage.size(); ++stage) {
        if (!has_ref[stage]) {
            continue;
        }
        if (std::norm(ref_by_stage[stage]) <= kReferenceMagnitudeSquaredEpsilon) {
            continue;
        }

        if (has_s11_measured[stage]) {
            const auto s11_ratio = s11_measured_by_stage[stage] / ref_by_stage[stage];
            datapoint_out.s11.re = s11_ratio.real();
            datapoint_out.s11.im = s11_ratio.imag();
            datapoint_out.has_s11 = true;
        }
        if (has_s21_measured[stage]) {
            const auto s21_ratio = s21_measured_by_stage[stage] / ref_by_stage[stage];
            datapoint_out.s21.re = s21_ratio.real();
            datapoint_out.s21.im = s21_ratio.imag();
            datapoint_out.has_s21 = true;
        }
        if (datapoint_out.has_s21) {
            return true;
        }
    }

    return false;
}

auto LibreVnaMinimalDriver::crc32(std::span<const std::uint8_t> data) -> std::uint32_t {
    std::uint32_t crc = 0xFFFFFFFFU;
    for (const auto byte : data) {
        crc ^= static_cast<std::uint32_t>(byte);
        for (int bit = 0; bit < 8; ++bit) {
            const auto lsb = crc & 1U;
            crc >>= 1U;
            if (lsb != 0U) {
                crc ^= 0xEDB88320U;
            }
        }
    }
    return ~crc;
}

}  // namespace radar::drivers
