#include "../compact_m_k209_driver.hpp"

#include <algorithm>
#include <array>
#include <cctype>
#include <cmath>
#include <cstring>
#include <limits>
#include <sstream>
#include <stdexcept>
#include <utility>

namespace radar::drivers {
namespace {

constexpr std::uint32_t kFloatBytes = 4U;
constexpr std::uint32_t kComplexScalarCount = 2U;
constexpr std::uint64_t kMaxBinaryBlockBytes = 512ULL * 1024ULL * 1024ULL;

[[nodiscard]] auto visa_resource(const std::string& resource) -> ViRsrc {
    return reinterpret_cast<ViRsrc>(const_cast<char*>(resource.c_str()));
}

[[nodiscard]] auto visa_buffer(std::uint8_t* data) -> ViBuf {
    return reinterpret_cast<ViBuf>(data);
}

[[nodiscard]] auto visa_const_buffer(const char* data) -> ViBuf {
    return reinterpret_cast<ViBuf>(const_cast<char*>(data));
}

[[nodiscard]] auto is_line_ending(std::uint8_t value) -> bool {
    return value == static_cast<std::uint8_t>('\n') || value == static_cast<std::uint8_t>('\r');
}

[[nodiscard]] auto is_response_separator(std::uint8_t value) -> bool {
    return is_line_ending(value) || value == static_cast<std::uint8_t>(';');
}

[[nodiscard]] auto checked_vi_count(std::uint64_t value) -> ViUInt32 {
    const auto max_count = static_cast<std::uint64_t>(std::numeric_limits<ViUInt32>::max());
    return static_cast<ViUInt32>(std::min(value, max_count));
}

}  // namespace

CompactMK209Driver::CompactMK209Driver(CompactMK209DriverSettings settings) : settings_(std::move(settings)) {}

CompactMK209Driver::~CompactMK209Driver() {
    try {
        close();
    } catch (...) {
    }
}

void CompactMK209Driver::open() {
    if (is_open_) {
        return;
    }
    if (settings_.resource.empty()) {
        throw std::runtime_error("K209 VISA resource must not be empty");
    }
    if (settings_.sweep.points < 2U) {
        throw std::runtime_error("K209 sweep points must be >= 2");
    }
    if (settings_.sweep.stop_hz < settings_.sweep.start_hz) {
        throw std::runtime_error("K209 sweep stop_hz must be >= start_hz");
    }
    if (!(settings_.sweep.if_bandwidth_hz > 0.0F)) {
        throw std::runtime_error("K209 IF bandwidth must be > 0");
    }

    ViStatus status = viOpenDefaultRM(&resource_manager_);
    if (status < VI_SUCCESS) {
        throw std::runtime_error("Failed to open VISA resource manager: status=" + std::to_string(status));
    }

    try {
        check_status(
            viOpen(resource_manager_, visa_resource(settings_.resource), VI_NULL, VI_NULL, &instrument_),
            "open K209 VISA resource"
        );
        check_status(
            viSetAttribute(instrument_, VI_ATTR_TMO_VALUE, static_cast<ViAttrState>(settings_.timeout_ms)),
            "set K209 VISA timeout"
        );
        check_status(
            viSetAttribute(instrument_, VI_ATTR_TERMCHAR_EN, static_cast<ViAttrState>(VI_FALSE)),
            "disable K209 VISA termchar"
        );
        write_command("*CLS");
        configure_device();
    } catch (...) {
        close();
        throw;
    }

    is_open_ = true;
}

void CompactMK209Driver::close() {
    if (instrument_ != VI_NULL) {
        viClose(instrument_);
        instrument_ = VI_NULL;
    }
    if (resource_manager_ != VI_NULL) {
        viClose(resource_manager_);
        resource_manager_ = VI_NULL;
    }
    is_open_ = false;
}

auto CompactMK209Driver::acquire_sweep() -> SweepTrace {
    require_open();

    auto traces = query_sweep_trace_pair(settings_.sweep.points);

    SweepTrace trace{};
    trace.frequency_hz = frequency_hz_;
    trace.s11 = complex_trace_from_interleaved(traces.first);
    trace.s21 = complex_trace_from_interleaved(traces.second);
    return trace;
}

auto CompactMK209Driver::query_identity() -> std::string {
    require_open();
    return trim_ascii(query_string("*IDN?"));
}

auto CompactMK209Driver::query_system_error() -> std::string {
    require_open();
    return trim_ascii(query_string("SYST:ERR?"));
}

void CompactMK209Driver::configure_device() {
    if (settings_.preset_on_open) {
        write_command("SYST:PRES");
        expect_operation_complete("*OPC?", "wait for K209 preset");
    }

    write_command("SENS:FREQ:STAR " + format_frequency(settings_.sweep.start_hz));
    write_command("SENS:FREQ:STOP " + format_frequency(settings_.sweep.stop_hz));
    write_command("SENS:SWE:POIN " + std::to_string(settings_.sweep.points));
    write_command("SENS:SWE:POIN:TIME 0");
    write_command("SENS:BAND " + format_frequency(settings_.sweep.if_bandwidth_hz));
    write_command("SOUR:POW " + format_power(settings_.sweep.power_dbm));
    write_command("SENS:AVER OFF");
    write_command("CALC:PAR1:DEF S21");
    write_command("CALC:PAR1:SEL");
    write_command("FORM:DATA REAL32");
    write_command("FORM:BORD SWAP");
    write_command("INIT:CONT ON");
    write_command("TRIG:SOUR BUS");
    expect_operation_complete("*OPC?", "wait for K209 setup");

    frequency_hz_ = query_float_array("SENS:FREQ:DATA?", settings_.sweep.points);
}

void CompactMK209Driver::expect_operation_complete(const std::string& command, const std::string& context) {
    const auto response = trim_ascii(query_string(command));
    if (response != "1") {
        throw std::runtime_error(context + " returned unexpected *OPC? response: " + response);
    }
}

void CompactMK209Driver::write_command(const std::string& command) {
    require_open();

    const std::string message = command + "\n";
    ViUInt32 written = 0;
    check_status(
        viWrite(instrument_, visa_const_buffer(message.data()), static_cast<ViUInt32>(message.size()), &written),
        "write K209 command: " + command
    );
    if (written != message.size()) {
        throw std::runtime_error("Incomplete K209 command write: " + command);
    }
}

auto CompactMK209Driver::query_string(const std::string& command) -> std::string {
    write_command(command);
    return read_line();
}

auto CompactMK209Driver::query_float_array(
    const std::string& command,
    std::uint32_t expected_values
) -> std::vector<float> {
    write_command(command);
    return read_float_array_response(command, expected_values);
}

auto CompactMK209Driver::query_sweep_trace_pair(std::uint32_t points)
    -> std::pair<std::vector<float>, std::vector<float>> {
    write_command("TRIG:SING;*OPC?;:SENS:DATA:CORR? S11;:SENS:DATA:CORR? S21");
    const auto opc_response = read_ascii_token();
    if (opc_response != "1") {
        throw std::runtime_error("K209 sweep returned unexpected *OPC? response: " + opc_response);
    }

    return {
        read_float_array_response("SENS:DATA:CORR? S11", points * kComplexScalarCount),
        read_float_array_response("SENS:DATA:CORR? S21", points * kComplexScalarCount),
    };
}

auto CompactMK209Driver::read_float_array_response(
    const std::string& context,
    std::uint32_t expected_values
) -> std::vector<float> {
    auto values = float_array_from_little_endian(read_ieee_block());
    if (values.size() != expected_values) {
        throw std::runtime_error(
            "K209 response for " + context + " returned " + std::to_string(values.size()) +
            " float32 values, expected " + std::to_string(expected_values)
        );
    }
    return values;
}

auto CompactMK209Driver::read_ascii_token() -> std::string {
    require_open();

    std::string token{};
    token.reserve(32);

    while (true) {
        const auto byte = read_byte();
        if (is_response_separator(byte)) {
            if (!token.empty()) {
                return token;
            }
            continue;
        }
        token.push_back(static_cast<char>(byte));
        if (token.size() > 64U * 1024U) {
            throw std::runtime_error("K209 ASCII response token is too long");
        }
    }
}

auto CompactMK209Driver::read_line() -> std::string {
    require_open();

    std::string line{};
    line.reserve(128);
    bool has_content = false;

    while (true) {
        const auto byte = read_byte();
        if (!has_content && is_line_ending(byte)) {
            continue;
        }
        if (byte == static_cast<std::uint8_t>('\n')) {
            break;
        }
        if (byte != static_cast<std::uint8_t>('\r')) {
            line.push_back(static_cast<char>(byte));
            has_content = true;
        }
        if (line.size() > 64U * 1024U) {
            throw std::runtime_error("K209 ASCII response line is too long");
        }
    }

    return line;
}

auto CompactMK209Driver::read_byte() -> std::uint8_t {
    auto bytes = read_exact(1);
    return bytes[0];
}

auto CompactMK209Driver::read_exact(std::uint64_t size) -> std::vector<std::uint8_t> {
    require_open();

    std::vector<std::uint8_t> bytes(size);
    std::uint64_t offset = 0;
    while (offset < size) {
        ViUInt32 transferred = 0;
        const auto count = checked_vi_count(size - offset);
        check_status(
            viRead(instrument_, visa_buffer(bytes.data() + offset), count, &transferred),
            "read K209 response bytes"
        );
        if (transferred == 0U) {
            throw std::runtime_error("K209 VISA read returned zero bytes before expected payload was complete");
        }
        offset += transferred;
    }
    return bytes;
}

auto CompactMK209Driver::read_ieee_block() -> std::vector<std::uint8_t> {
    std::uint8_t marker = read_byte();
    while (is_response_separator(marker)) {
        marker = read_byte();
    }
    if (marker != static_cast<std::uint8_t>('#')) {
        throw std::runtime_error("K209 binary response does not start with IEEE block marker '#'");
    }

    const auto digit = read_byte();
    if (digit != static_cast<std::uint8_t>('8')) {
        throw std::runtime_error("K209 binary response uses unsupported IEEE block header width");
    }

    const auto size_text_bytes = read_exact(8);
    std::uint64_t payload_size = 0;
    for (const auto byte : size_text_bytes) {
        if (!std::isdigit(static_cast<unsigned char>(byte))) {
            throw std::runtime_error("K209 binary response has non-numeric payload size");
        }
        payload_size = (payload_size * 10ULL) + static_cast<std::uint64_t>(byte - static_cast<std::uint8_t>('0'));
    }
    if (payload_size == 0U || payload_size > kMaxBinaryBlockBytes) {
        throw std::runtime_error("K209 binary response payload size is out of allowed range");
    }
    if ((payload_size % kFloatBytes) != 0U) {
        throw std::runtime_error("K209 binary response payload size is not aligned to float32 values");
    }

    return read_exact(payload_size);
}

void CompactMK209Driver::check_status(ViStatus status, const std::string& context) const {
    if (status >= VI_SUCCESS) {
        return;
    }
    throw std::runtime_error(context + ": " + status_message(status));
}

auto CompactMK209Driver::status_message(ViStatus status) const -> std::string {
    std::array<ViChar, 256> description{};
    const auto session = instrument_ != VI_NULL ? instrument_ : resource_manager_;
    if (session != VI_NULL && viStatusDesc(session, status, description.data()) >= VI_SUCCESS) {
        return std::string(description.data()) + " (status=" + std::to_string(status) + ")";
    }
    return "VISA status=" + std::to_string(status);
}

void CompactMK209Driver::require_open() const {
    if (instrument_ == VI_NULL) {
        throw std::runtime_error("K209 VISA instrument is not open");
    }
}

auto CompactMK209Driver::float_array_from_little_endian(std::vector<std::uint8_t> bytes) -> std::vector<float> {
    if ((bytes.size() % kFloatBytes) != 0U) {
        throw std::runtime_error("K209 binary float payload is not aligned to 4 bytes");
    }

    std::vector<float> values(bytes.size() / kFloatBytes, 0.0F);
    for (std::size_t index = 0; index < values.size(); ++index) {
        const auto offset = index * kFloatBytes;
        const std::uint32_t bits = static_cast<std::uint32_t>(bytes[offset]) |
                                   (static_cast<std::uint32_t>(bytes[offset + 1U]) << 8U) |
                                   (static_cast<std::uint32_t>(bytes[offset + 2U]) << 16U) |
                                   (static_cast<std::uint32_t>(bytes[offset + 3U]) << 24U);
        std::memcpy(&values[index], &bits, sizeof(float));
    }
    return values;
}

auto CompactMK209Driver::complex_trace_from_interleaved(const std::vector<float>& values)
    -> std::vector<ipc::Complex32> {
    if ((values.size() % kComplexScalarCount) != 0U) {
        throw std::runtime_error("K209 complex trace payload has odd scalar count");
    }

    std::vector<ipc::Complex32> trace{};
    trace.reserve(values.size() / kComplexScalarCount);
    for (std::size_t index = 0; index < values.size(); index += kComplexScalarCount) {
        trace.push_back(ipc::Complex32{.re = values[index], .im = values[index + 1U]});
    }
    return trace;
}

auto CompactMK209Driver::format_frequency(float value_hz) -> std::string {
    if (!std::isfinite(value_hz)) {
        throw std::runtime_error("K209 frequency value is not finite");
    }

    std::ostringstream stream;
    stream.precision(9);
    stream << std::fixed << value_hz;
    return stream.str();
}

auto CompactMK209Driver::format_power(float value_dbm) -> std::string {
    if (!std::isfinite(value_dbm)) {
        throw std::runtime_error("K209 power value is not finite");
    }

    std::ostringstream stream;
    stream.precision(3);
    stream << std::fixed << value_dbm;
    return stream.str();
}

auto CompactMK209Driver::trim_ascii(std::string value) -> std::string {
    const auto first = std::find_if_not(value.begin(), value.end(), [](unsigned char ch) {
        return std::isspace(ch) != 0;
    });
    const auto last = std::find_if_not(value.rbegin(), value.rend(), [](unsigned char ch) {
        return std::isspace(ch) != 0;
    }).base();
    if (first >= last) {
        return {};
    }
    return std::string(first, last);
}

}  // namespace radar::drivers
