#pragma once

#include <cstdint>
#include <span>
#include <vector>

#include "wireframe/net/ipv4.hpp"

// RFC 1071 Internet checksum, and the IPv4/TCP/UDP verification built
// on it. Not wired into summarize_packet(): on loopback, and for many
// packets captured right as they leave the local machine, the
// transmitted checksum is legitimately 0x0000 or garbage - modern
// NICs compute it in hardware ("checksum offload") only once the frame
// actually reaches them, which is *after* most capture points see it.
// Flagging that as "BAD" by default would be noise, not signal, on
// exactly the interfaces this project has been tested against all
// session (lo, tailscale0). Wireshark makes this opt-in for the same
// reason; so does this (CLI's -c flag calls these directly).
namespace wireframe::net {

// One's-complement sum of 16-bit big-endian words, folded back into 16
// bits, then complemented. Used identically by IPv4's header checksum
// and, over a pseudo-header + segment instead of a plain header, by
// TCP/UDP.
inline std::uint16_t internet_checksum(std::span<const unsigned char> data) {
    std::uint32_t sum = 0;
    std::size_t i = 0;
    for (; i + 1 < data.size(); i += 2) {
        sum += (static_cast<std::uint32_t>(data[i]) << 8) | data[i + 1];
    }
    if (i < data.size()) {
        sum += static_cast<std::uint32_t>(data[i]) << 8;  // odd trailing byte: high half only
    }
    while (sum >> 16) {
        sum = (sum & 0xFFFFu) + (sum >> 16);
    }
    return static_cast<std::uint16_t>(~sum & 0xFFFFu);
}

// `header_bytes` must be exactly the IPv4 header as it appeared on the
// wire (IHL*4 bytes, options included, checksum field included as its
// real transmitted value - not zeroed). Summing a header that already
// contains its own valid checksum comes out to exactly 0; that's the
// verification, no need for a mutable copy with the field zeroed out.
inline bool verify_ipv4_checksum(std::span<const unsigned char> header_bytes) {
    return internet_checksum(header_bytes) == 0;
}

enum class ChecksumResult { kValid, kInvalid, kNotPresent };

namespace detail {

inline std::vector<unsigned char> build_ipv4_pseudo_header(const Ipv4Address& src,
                                                             const Ipv4Address& dst,
                                                             std::uint8_t protocol,
                                                             std::span<const unsigned char> segment) {
    std::vector<unsigned char> buf;
    buf.reserve(12 + segment.size());
    buf.insert(buf.end(), src.bytes.begin(), src.bytes.end());
    buf.insert(buf.end(), dst.bytes.begin(), dst.bytes.end());
    buf.push_back(0);
    buf.push_back(protocol);
    std::uint16_t len = static_cast<std::uint16_t>(segment.size());
    buf.push_back(static_cast<unsigned char>(len >> 8));
    buf.push_back(static_cast<unsigned char>(len & 0xFF));
    buf.insert(buf.end(), segment.begin(), segment.end());
    return buf;
}

}  // namespace detail

// TCP's checksum is mandatory - always kValid or kInvalid.
inline ChecksumResult verify_tcp_checksum_ipv4(const Ipv4Address& src, const Ipv4Address& dst,
                                                std::span<const unsigned char> tcp_segment) {
    auto buf = detail::build_ipv4_pseudo_header(src, dst, kProtoTcp, tcp_segment);
    return internet_checksum(buf) == 0 ? ChecksumResult::kValid : ChecksumResult::kInvalid;
}

// UDP's checksum is optional over IPv4 (RFC 768): a transmitted value
// of exactly 0x0000 means "no checksum was computed", not "checksum is
// zero" - that's kNotPresent, not a failure.
inline ChecksumResult verify_udp_checksum_ipv4(const Ipv4Address& src, const Ipv4Address& dst,
                                                std::span<const unsigned char> udp_datagram) {
    if (udp_datagram.size() >= 8 && udp_datagram[6] == 0 && udp_datagram[7] == 0) {
        return ChecksumResult::kNotPresent;
    }
    auto buf = detail::build_ipv4_pseudo_header(src, dst, kProtoUdp, udp_datagram);
    return internet_checksum(buf) == 0 ? ChecksumResult::kValid : ChecksumResult::kInvalid;
}

}  // namespace wireframe::net