#include "shared_types.hpp"

#include <chrono>
#include <cstring>
#include <limits>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>

namespace radar::ipc {
namespace {

constexpr std::uint32_t kRawCollectionMagic = 0x32574152U;      // RAW2
constexpr std::uint32_t kPreprocessedCollectionMagic = 0x32525050U;  // PRP2
constexpr std::uint32_t kResultCollectionMagic = 0x314C5352U;   // RSL1

template <typename T>
concept TriviallySerializable = std::is_trivially_copyable_v<T>;

[[nodiscard]] auto checked_count_to_u32(std::size_t count, const std::string& label) -> std::uint32_t {
    if (count > std::numeric_limits<std::uint32_t>::max()) {
        throw std::runtime_error(label + " exceeds uint32 wire-format limit");
    }
    return static_cast<std::uint32_t>(count);
}

void ensure_equal_sizes(std::size_t left, std::size_t right, const std::string& label) {
    if (left != right) {
        throw std::runtime_error(label + " has inconsistent vector sizes");
    }
}

class BinaryWriter {
  public:
    template <TriviallySerializable T>
    void write(const T& value) {
        const auto old_size = bytes_.size();
        bytes_.resize(old_size + sizeof(T));
        std::memcpy(bytes_.data() + old_size, &value, sizeof(T));
    }

    void write_string(const std::string& value) {
        if (value.size() > std::numeric_limits<std::uint16_t>::max()) {
            throw std::runtime_error("String is too large for wire format");
        }

        write(static_cast<std::uint16_t>(value.size()));

        const auto old_size = bytes_.size();
        bytes_.resize(old_size + value.size());
        std::memcpy(bytes_.data() + old_size, value.data(), value.size());
    }

    [[nodiscard]] auto finish() && -> std::vector<std::uint8_t> {
        return std::move(bytes_);
    }

  private:
    std::vector<std::uint8_t> bytes_{};
};

class BinaryReader {
  public:
    explicit BinaryReader(std::span<const std::uint8_t> bytes) : bytes_(bytes) {}

    template <TriviallySerializable T>
    [[nodiscard]] auto read() -> T {
        ensure_available(sizeof(T));

        T value{};
        std::memcpy(&value, bytes_.data() + offset_, sizeof(T));
        offset_ += sizeof(T);
        return value;
    }

    [[nodiscard]] auto read_string() -> std::string {
        const auto size = read<std::uint16_t>();
        ensure_available(size);

        const auto* begin = reinterpret_cast<const char*>(bytes_.data() + offset_);
        std::string value(begin, begin + size);
        offset_ += size;
        return value;
    }

    [[nodiscard]] auto is_consumed() const -> bool {
        return offset_ == bytes_.size();
    }

    [[nodiscard]] auto remaining_bytes() const -> std::size_t {
        return bytes_.size() - offset_;
    }

  private:
    void ensure_available(std::size_t size) const {
        if (offset_ + size > bytes_.size()) {
            throw std::runtime_error("Unexpected end of serialized payload");
        }
    }

    std::span<const std::uint8_t> bytes_{};
    std::size_t offset_ = 0;
};

void write_trace_block(BinaryWriter& writer, const SweepTraceBlock& trace) {
    ensure_equal_sizes(trace.frequency_hz.size(), trace.s21.size(), "Sweep trace");
    ensure_equal_sizes(trace.frequency_hz.size(), trace.s11.size(), "Sweep trace S11");

    writer.write(trace.combo.input_pos);
    writer.write(trace.combo.output_pos);
    writer.write(checked_count_to_u32(trace.frequency_hz.size(), "Trace point count"));

    for (const auto frequency_hz : trace.frequency_hz) {
        writer.write(frequency_hz);
    }

    for (const auto& point : trace.s11) {
        writer.write(point.re);
        writer.write(point.im);
    }

    for (const auto& point : trace.s21) {
        writer.write(point.re);
        writer.write(point.im);
    }
}

[[nodiscard]] auto read_trace_block(BinaryReader& reader) -> SweepTraceBlock {
    SweepTraceBlock trace{};
    trace.combo.input_pos = reader.read<std::uint32_t>();
    trace.combo.output_pos = reader.read<std::uint32_t>();

    const auto point_count = reader.read<std::uint32_t>();
    trace.frequency_hz.reserve(point_count);
    trace.s11.reserve(point_count);
    trace.s21.reserve(point_count);

    for (std::uint32_t index = 0; index < point_count; ++index) {
        trace.frequency_hz.push_back(reader.read<float>());
    }

    for (std::uint32_t index = 0; index < point_count; ++index) {
        trace.s11.push_back(Complex32{
            .re = reader.read<float>(),
            .im = reader.read<float>(),
        });
    }

    for (std::uint32_t index = 0; index < point_count; ++index) {
        trace.s21.push_back(Complex32{
            .re = reader.read<float>(),
            .im = reader.read<float>(),
        });
    }

    return trace;
}

void write_trace_collection(BinaryWriter& writer, std::uint32_t magic, const RawSweepCollection& collection) {
    writer.write(magic);
    writer.write(collection.collection_id);
    writer.write(collection.monotonic_ns);
    writer.write(checked_count_to_u32(collection.traces.size(), "Trace count"));

    for (const auto& trace : collection.traces) {
        write_trace_block(writer, trace);
    }

    writer.write(collection.capture_start_ns);
    writer.write(collection.capture_end_ns);
}

[[nodiscard]] auto read_trace_collection(BinaryReader& reader, std::uint32_t expected_magic) -> RawSweepCollection {
    const auto magic = reader.read<std::uint32_t>();
    if (magic != expected_magic) {
        throw std::runtime_error("Unexpected trace collection magic");
    }

    RawSweepCollection collection{};
    collection.collection_id = reader.read<std::uint64_t>();
    collection.monotonic_ns = reader.read<std::uint64_t>();

    const auto trace_count = reader.read<std::uint32_t>();
    collection.traces.reserve(trace_count);
    for (std::uint32_t index = 0; index < trace_count; ++index) {
        collection.traces.push_back(read_trace_block(reader));
    }

    if (reader.remaining_bytes() == 0U) {
        return collection;
    }
    if (reader.remaining_bytes() != (sizeof(std::uint64_t) * 2U)) {
        throw std::runtime_error("Unexpected trailing bytes in trace collection");
    }

    collection.capture_start_ns = reader.read<std::uint64_t>();
    collection.capture_end_ns = reader.read<std::uint64_t>();

    return collection;
}

void write_trace_result_payload(BinaryWriter& writer, const ResultPayload& payload) {
    ensure_equal_sizes(payload.frequency_hz.size(), payload.trace.size(), "Result trace payload");

    writer.write(checked_count_to_u32(payload.trace.size(), "Result trace point count"));
    for (const auto frequency_hz : payload.frequency_hz) {
        writer.write(frequency_hz);
    }
    for (const auto& sample : payload.trace) {
        writer.write(sample.re);
        writer.write(sample.im);
    }
}

auto read_trace_result_payload(BinaryReader& reader, ResultPayload* payload) -> void {
    const auto point_count = reader.read<std::uint32_t>();

    payload->frequency_hz.reserve(point_count);
    payload->trace.reserve(point_count);

    for (std::uint32_t index = 0; index < point_count; ++index) {
        payload->frequency_hz.push_back(reader.read<float>());
    }
    for (std::uint32_t index = 0; index < point_count; ++index) {
        payload->trace.push_back(Complex32{
            .re = reader.read<float>(),
            .im = reader.read<float>(),
        });
    }
}

void write_image_result_payload(BinaryWriter& writer, const ResultPayload& payload) {
    const auto x_count = checked_count_to_u32(payload.image_x_axis.size(), "Result image x axis point count");
    const auto y_count = checked_count_to_u32(payload.image_y_axis.size(), "Result image y axis point count");
    const std::size_t expected_value_count = static_cast<std::size_t>(x_count) * static_cast<std::size_t>(y_count);
    if (payload.image_values.size() != expected_value_count) {
        throw std::runtime_error("Result image payload has inconsistent axis/value sizes");
    }

    writer.write(x_count);
    writer.write(y_count);
    for (const auto value : payload.image_x_axis) {
        writer.write(value);
    }
    for (const auto value : payload.image_y_axis) {
        writer.write(value);
    }
    for (const auto value : payload.image_values) {
        writer.write(value);
    }
}

auto read_image_result_payload(BinaryReader& reader, ResultPayload* payload) -> void {
    const auto x_count = reader.read<std::uint32_t>();
    const auto y_count = reader.read<std::uint32_t>();

    payload->image_x_axis.reserve(x_count);
    payload->image_y_axis.reserve(y_count);
    payload->image_values.reserve(static_cast<std::size_t>(x_count) * static_cast<std::size_t>(y_count));

    for (std::uint32_t index = 0; index < x_count; ++index) {
        payload->image_x_axis.push_back(reader.read<float>());
    }
    for (std::uint32_t index = 0; index < y_count; ++index) {
        payload->image_y_axis.push_back(reader.read<float>());
    }
    for (std::size_t index = 0; index < static_cast<std::size_t>(x_count) * static_cast<std::size_t>(y_count); ++index) {
        payload->image_values.push_back(reader.read<float>());
    }
}

void write_table_result_payload(BinaryWriter& writer, const ResultPayload& payload) {
    const auto column_count = payload.table_columns;
    if (column_count == 0U && !payload.table_values.empty()) {
        throw std::runtime_error("Result table payload must declare table_columns when values are present");
    }
    if (column_count != 0U && (payload.table_values.size() % column_count) != 0U) {
        throw std::runtime_error("Result table payload has inconsistent column/value sizes");
    }

    const std::uint32_t row_count =
        column_count == 0U ? 0U : checked_count_to_u32(payload.table_values.size() / column_count, "Result table row count");

    writer.write(column_count);
    writer.write(row_count);
    for (const auto value : payload.table_values) {
        writer.write(value);
    }
}

auto read_table_result_payload(BinaryReader& reader, ResultPayload* payload) -> void {
    const auto column_count = reader.read<std::uint32_t>();
    const auto row_count = reader.read<std::uint32_t>();
    payload->table_columns = column_count;
    payload->table_values.reserve(static_cast<std::size_t>(column_count) * static_cast<std::size_t>(row_count));
    for (std::size_t index = 0; index < static_cast<std::size_t>(column_count) * static_cast<std::size_t>(row_count); ++index) {
        payload->table_values.push_back(reader.read<float>());
    }
}

void write_result_payload(BinaryWriter& writer, const ResultPayload& payload) {
    writer.write(static_cast<std::uint8_t>(payload.kind));
    writer.write_string(payload.processing_name);

    switch (payload.kind) {
        case ResultKind::TraceComplex:
            write_trace_result_payload(writer, payload);
            return;
        case ResultKind::ScalarF32:
            writer.write(payload.scalar_value);
            return;
        case ResultKind::ImageF32:
            write_image_result_payload(writer, payload);
            return;
        case ResultKind::TableF32:
            write_table_result_payload(writer, payload);
            return;
        default:
            throw std::runtime_error("Unsupported result payload kind");
    }
}

[[nodiscard]] auto read_result_payload(BinaryReader& reader) -> ResultPayload {
    ResultPayload payload{};
    payload.kind = static_cast<ResultKind>(reader.read<std::uint8_t>());
    payload.processing_name = reader.read_string();

    switch (payload.kind) {
        case ResultKind::TraceComplex:
            read_trace_result_payload(reader, &payload);
            return payload;
        case ResultKind::ScalarF32:
            payload.scalar_value = reader.read<float>();
            return payload;
        case ResultKind::ImageF32:
            read_image_result_payload(reader, &payload);
            return payload;
        case ResultKind::TableF32:
            read_table_result_payload(reader, &payload);
            return payload;
        default:
            throw std::runtime_error("Unsupported result payload kind in stream");
    }
}

void write_result_block(BinaryWriter& writer, const ResultBlock& block) {
    writer.write(block.combo.input_pos);
    writer.write(block.combo.output_pos);
    writer.write(checked_count_to_u32(block.payloads.size(), "Result payload count"));

    for (const auto& payload : block.payloads) {
        write_result_payload(writer, payload);
    }
}

[[nodiscard]] auto read_result_block(BinaryReader& reader) -> ResultBlock {
    ResultBlock block{};
    block.combo.input_pos = reader.read<std::uint32_t>();
    block.combo.output_pos = reader.read<std::uint32_t>();

    const auto payload_count = reader.read<std::uint32_t>();
    block.payloads.reserve(payload_count);
    for (std::uint32_t index = 0; index < payload_count; ++index) {
        block.payloads.push_back(read_result_payload(reader));
    }

    return block;
}

void ensure_reader_consumed(const BinaryReader& reader, const std::string& payload_label) {
    if (!reader.is_consumed()) {
        throw std::runtime_error("Unexpected trailing bytes in " + payload_label);
    }
}

}  // namespace

auto serialize_raw_collection(const RawSweepCollection& collection) -> std::vector<std::uint8_t> {
    BinaryWriter writer{};
    write_trace_collection(writer, kRawCollectionMagic, collection);
    return std::move(writer).finish();
}

auto deserialize_raw_collection(std::span<const std::uint8_t> bytes) -> RawSweepCollection {
    BinaryReader reader(bytes);
    auto collection = read_trace_collection(reader, kRawCollectionMagic);
    ensure_reader_consumed(reader, "raw collection");
    return collection;
}

auto serialize_preprocessed_collection(const PreprocessedCollection& collection) -> std::vector<std::uint8_t> {
    BinaryWriter writer{};
    write_trace_collection(writer, kPreprocessedCollectionMagic, collection);
    return std::move(writer).finish();
}

auto deserialize_preprocessed_collection(std::span<const std::uint8_t> bytes) -> PreprocessedCollection {
    BinaryReader reader(bytes);
    auto collection = read_trace_collection(reader, kPreprocessedCollectionMagic);
    ensure_reader_consumed(reader, "preprocessed collection");
    return collection;
}

auto serialize_result_collection(const ResultCollection& collection) -> std::vector<std::uint8_t> {
    BinaryWriter writer{};
    writer.write(kResultCollectionMagic);
    writer.write(collection.collection_id);
    writer.write(collection.monotonic_ns);
    writer.write(checked_count_to_u32(collection.collection_payloads.size(), "Collection payload count"));
    writer.write(checked_count_to_u32(collection.blocks.size(), "Result block count"));

    for (const auto& payload : collection.collection_payloads) {
        write_result_payload(writer, payload);
    }
    for (const auto& block : collection.blocks) {
        write_result_block(writer, block);
    }

    return std::move(writer).finish();
}

auto deserialize_result_collection(std::span<const std::uint8_t> bytes) -> ResultCollection {
    BinaryReader reader(bytes);

    const auto magic = reader.read<std::uint32_t>();
    if (magic != kResultCollectionMagic) {
        throw std::runtime_error("Unexpected result collection magic");
    }

    ResultCollection collection{};
    collection.collection_id = reader.read<std::uint64_t>();
    collection.monotonic_ns = reader.read<std::uint64_t>();

    const auto collection_payload_count = reader.read<std::uint32_t>();
    const auto block_count = reader.read<std::uint32_t>();
    collection.collection_payloads.reserve(collection_payload_count);
    for (std::uint32_t index = 0; index < collection_payload_count; ++index) {
        collection.collection_payloads.push_back(read_result_payload(reader));
    }

    collection.blocks.reserve(block_count);
    for (std::uint32_t index = 0; index < block_count; ++index) {
        collection.blocks.push_back(read_result_block(reader));
    }

    ensure_reader_consumed(reader, "result collection");
    return collection;
}

auto current_monotonic_ns() -> std::uint64_t {
    const auto now = std::chrono::steady_clock::now().time_since_epoch();
    const auto now_ns = std::chrono::duration_cast<std::chrono::nanoseconds>(now).count();
    return static_cast<std::uint64_t>(now_ns);
}

}  // namespace radar::ipc
