#include <doctest/doctest.h>

#include <vector>

#include "packeteer/l7/snmp.hpp"

using namespace packeteer::net;

namespace {

std::vector<unsigned char> ber_length(std::size_t n) {
    if (n < 0x80) return {static_cast<unsigned char>(n)};
    std::vector<unsigned char> bytes;
    while (n) {
        bytes.insert(bytes.begin(), static_cast<unsigned char>(n & 0xFF));
        n >>= 8;
    }
    bytes.insert(bytes.begin(), static_cast<unsigned char>(0x80 | bytes.size()));
    return bytes;
}

std::vector<unsigned char> ber_tlv(std::uint8_t tag, const std::vector<unsigned char>& value) {
    std::vector<unsigned char> out = {tag};
    auto len = ber_length(value.size());
    out.insert(out.end(), len.begin(), len.end());
    out.insert(out.end(), value.begin(), value.end());
    return out;
}

std::vector<unsigned char> ber_int(std::int64_t n) {
    if (n == 0) return {0x00};
    std::vector<unsigned char> bytes;
    bool neg = n < 0;
    std::uint64_t val = neg ? static_cast<std::uint64_t>(~n) : static_cast<std::uint64_t>(n);
    while (val) {
        bytes.insert(bytes.begin(), static_cast<unsigned char>(val & 0xFF));
        val >>= 8;
    }
    if (!neg && (bytes[0] & 0x80)) bytes.insert(bytes.begin(), 0x00);
    return bytes;
}

std::vector<unsigned char> ber_octet_string(const std::string& s) {
    return std::vector<unsigned char>(s.begin(), s.end());
}

std::vector<unsigned char> concat(std::initializer_list<std::vector<unsigned char>> parts) {
    std::vector<unsigned char> out;
    for (const auto& p : parts) out.insert(out.end(), p.begin(), p.end());
    return out;
}

std::vector<unsigned char> snmp_message(std::int64_t version, const std::string& community,
                                          std::uint8_t pdu_tag, std::int64_t request_id) {
    auto pdu_body =
        concat({ber_tlv(kBerTagInteger, ber_int(request_id)), ber_tlv(kBerTagInteger, ber_int(0)),
                ber_tlv(kBerTagInteger, ber_int(0)), ber_tlv(kBerTagSequence, {})});
    auto msg_body = concat({ber_tlv(kBerTagInteger, ber_int(version)),
                             ber_tlv(kBerTagOctetString, ber_octet_string(community)),
                             ber_tlv(pdu_tag, pdu_body)});
    return ber_tlv(kBerTagSequence, msg_body);
}

}  // namespace

TEST_CASE("parse_snmp decodes a v2c GetRequest") {
    auto bytes = snmp_message(1, "public", 0xA0, 12345);
    auto msg = parse_snmp(bytes);
    REQUIRE(msg.has_value());
    CHECK(msg->version == 1);
    REQUIRE(msg->community.has_value());
    CHECK(*msg->community == "public");
    REQUIRE(msg->pdu_tag.has_value());
    CHECK(*msg->pdu_tag == 0xA0);
    REQUIRE(msg->request_id.has_value());
    CHECK(*msg->request_id == 12345);
}

TEST_CASE("parse_snmp decodes a v1 message") {
    auto bytes = snmp_message(0, "private", 0xA2, 7);
    auto msg = parse_snmp(bytes);
    REQUIRE(msg.has_value());
    CHECK(msg->version == 0);
    REQUIRE(msg->community.has_value());
    CHECK(*msg->community == "private");
}

TEST_CASE("parse_snmp reports v3 by version alone, without community or pdu_tag") {
    // A minimal v3 SEQUENCE { version } - real v3 messages carry a
    // security-parameters header afterward, but version alone is
    // enough to trigger the "not decoded further" path.
    auto bytes = ber_tlv(kBerTagSequence, ber_tlv(kBerTagInteger, ber_int(3)));
    auto msg = parse_snmp(bytes);
    REQUIRE(msg.has_value());
    CHECK(msg->version == 3);
    CHECK_FALSE(msg->community.has_value());
    CHECK_FALSE(msg->pdu_tag.has_value());
}

TEST_CASE("parse_snmp rejects an unsupported version (e.g. the abandoned SNMPv2 variant)") {
    auto bytes = ber_tlv(kBerTagSequence, ber_tlv(kBerTagInteger, ber_int(2)));
    CHECK_FALSE(parse_snmp(bytes).has_value());
}

TEST_CASE("parse_snmp handles a v1 Trap-PDU without a request-id field") {
    // Trap-PDU's actual first field is an enterprise OID, not
    // request-id - deliberately don't build one, just confirm
    // request_id stays unset rather than misreading the OID as one.
    auto msg_body = concat({ber_tlv(kBerTagInteger, ber_int(0)),
                             ber_tlv(kBerTagOctetString, ber_octet_string("public")),
                             ber_tlv(kSnmpPduTrapV1, ber_tlv(0x06, {0x2b, 0x06, 0x01}))});
    auto bytes = ber_tlv(kBerTagSequence, msg_body);
    auto msg = parse_snmp(bytes);
    REQUIRE(msg.has_value());
    REQUIRE(msg->pdu_tag.has_value());
    CHECK(*msg->pdu_tag == kSnmpPduTrapV1);
    CHECK_FALSE(msg->request_id.has_value());
}

TEST_CASE("parse_snmp rejects a payload that isn't a BER SEQUENCE") {
    std::vector<unsigned char> bytes = {0x02, 0x01, 0x00};
    CHECK_FALSE(parse_snmp(bytes).has_value());
}

TEST_CASE("parse_snmp rejects a truncated message") {
    auto full = snmp_message(1, "public", 0xA0, 1);
    std::vector<unsigned char> truncated(full.begin(), full.begin() + 5);
    CHECK_FALSE(parse_snmp(truncated).has_value());
}

TEST_CASE("SnmpDissector claims the port given to its constructor") {
    SnmpDissector agent(kSnmpAgentPort);
    SnmpDissector trap(kSnmpTrapPort);
    CHECK(agent.port() == 161);
    CHECK(trap.port() == 162);
}

TEST_CASE("SnmpDissector formats a v2c GetRequest") {
    SnmpDissector dissector(kSnmpAgentPort);
    auto summary = dissector.summarize(snmp_message(1, "public", 0xA0, 12345));
    REQUIRE(summary.has_value());
    CHECK(*summary == "SNMP v2c GetRequest community=public request-id=12345");
}

TEST_CASE("SnmpDissector reports v3 without a community string") {
    SnmpDissector dissector(kSnmpAgentPort);
    auto bytes = ber_tlv(kBerTagSequence, ber_tlv(kBerTagInteger, ber_int(3)));
    auto summary = dissector.summarize(bytes);
    REQUIRE(summary.has_value());
    CHECK(*summary == "SNMP v3 (encrypted/authenticated, not decoded)");
}