#pragma once

#include <algorithm>
#include <array>
#include <cstdint>
#include <cstdio>
#include <optional>
#include <span>
#include <string>

#include "wireframe/byteio.hpp"

namespace wireframe::net {

inline constexpr std::size_t kIpv6HeaderLen = 40;
inline constexpr std::uint8_t kNextHeaderHopByHop = 0;
inline constexpr std::uint8_t kNextHeaderRouting = 43;
inline constexpr std::uint8_t kNextHeaderFragment = 44;
inline constexpr std::uint8_t kNextHeaderEsp = 50;
inline constexpr std::uint8_t kNextHeaderAh = 51;
inline constexpr std::uint8_t kNextHeaderIcmpv6 = 58;
inline constexpr std::uint8_t kNextHeaderDestOptions = 60;

struct Ipv6Address {
    std::array<unsigned char, 16> bytes;
};

struct Ipv6Header {
    std::uint8_t version;
    std::uint8_t traffic_class;
    std::uint32_t flow_label;
    std::uint16_t payload_length;
    std::uint8_t next_header;  // transport protocol, or an extension header type
    std::uint8_t hop_limit;
    Ipv6Address src;
    Ipv6Address dst;
};

struct Ipv6Packet {
    Ipv6Header header;
    std::span<const unsigned char> payload;
};

// Only the fixed 40-byte header is decoded here - header.next_header
// may name an extension header rather than a transport protocol.
// walk_ipv6_extension_headers() (below) resolves that; parse_ipv6()
// itself stays a direct, unconditional decode of exactly the fixed
// header, nothing more.
inline std::optional<Ipv6Packet> parse_ipv6(std::span<const unsigned char> bytes) {
    if (bytes.size() < kIpv6HeaderLen) return std::nullopt;

    std::uint8_t version = static_cast<std::uint8_t>(bytes[0] >> 4);
    if (version != 6) return std::nullopt;

    Ipv6Header header{};
    header.version = version;
    std::uint32_t first_word = read_be32(bytes, 0);
    header.traffic_class = static_cast<std::uint8_t>((first_word >> 20) & 0xFF);
    header.flow_label = first_word & 0x000FFFFF;
    header.payload_length = read_be16(bytes, 4);
    header.next_header = bytes[6];
    header.hop_limit = bytes[7];
    std::copy_n(bytes.begin() + 8, 16, header.src.bytes.begin());
    std::copy_n(bytes.begin() + 24, 16, header.dst.bytes.begin());

    return Ipv6Packet{header, bytes.subspan(kIpv6HeaderLen)};
}

struct Ipv6ExtensionWalkResult {
    std::uint8_t final_next_header;  // a transport protocol, or an extension type we stopped at
    std::span<const unsigned char> payload;  // bytes after every extension header walked
    bool stopped_at_esp;  // true if ESP was hit - see walk_ipv6_extension_headers()
};

// Walks Hop-by-Hop, Routing, Destination Options, Fragment, and AH
// extension headers to find the real transport protocol underneath
// them, so e.g. TCP/UDP wrapped in a Hop-by-Hop options header is still
// decoded instead of silently stopping at "next_header=0". Each header
// carries its own length, so this never needs to understand a header
// type's *meaning* to skip over it correctly - only Hop-by-Hop/
// Routing/Dest-Options (length in 8-byte units from a trailing byte),
// Fragment (fixed 8 bytes), and AH (length in 4-byte units, RFC 4302)
// have different encodings, all handled explicitly below.
//
// ESP is a hard stop, not a bug: its own next-header field lives in a
// trailer *after* the encrypted payload, at an offset this code has no
// way to know without decrypting first. Reported as stopped_at_esp
// rather than guessed at.
//
// Bounded to a handful of iterations as defense in depth against a
// hostile/corrupt chain - not strictly needed for termination (every
// header is at least 8 bytes, so payload.size() strictly decreases
// each iteration and the loop can't actually run forever), but a
// pathological chain of many tiny headers would otherwise still cost
// real work for no legitimate reason.
inline Ipv6ExtensionWalkResult walk_ipv6_extension_headers(std::uint8_t next_header,
                                                             std::span<const unsigned char> payload) {
    constexpr int kMaxExtensionHeaders = 8;

    for (int i = 0; i < kMaxExtensionHeaders; ++i) {
        if (next_header == kNextHeaderEsp) {
            return {next_header, payload, /*stopped_at_esp=*/true};
        }

        std::size_t ext_len;
        if (next_header == kNextHeaderFragment) {
            if (payload.size() < 8) return {next_header, payload, false};
            ext_len = 8;
        } else if (next_header == kNextHeaderAh) {
            if (payload.size() < 2) return {next_header, payload, false};
            ext_len = (static_cast<std::size_t>(payload[1]) + 2) * 4;
        } else if (next_header == kNextHeaderHopByHop || next_header == kNextHeaderRouting ||
                   next_header == kNextHeaderDestOptions) {
            if (payload.size() < 2) return {next_header, payload, false};
            ext_len = (static_cast<std::size_t>(payload[1]) + 1) * 8;
        } else {
            break;  // TCP/UDP/ICMPv6/anything else we don't chain through: stop here
        }

        if (payload.size() < ext_len) return {next_header, payload, false};  // truncated: stop

        std::uint8_t this_next_header = payload[0];
        payload = payload.subspan(ext_len);
        next_header = this_next_header;
    }

    return {next_header, payload, false};
}

// RFC 5952 canonical text form: lowercase hex, and the longest run of
// two-or-more consecutive zero groups (leftmost wins a tie) collapsed to
// "::". A lone zero group is left as "0", not compressed, per 5952 4.2.2.
inline std::string ipv6_to_string(const Ipv6Address& addr) {
    std::array<std::uint16_t, 8> groups{};
    for (std::size_t i = 0; i < 8; ++i) {
        groups[i] = static_cast<std::uint16_t>((addr.bytes[i * 2] << 8) | addr.bytes[i * 2 + 1]);
    }

    int best_start = -1;
    int best_len = 0;
    int cur_start = -1;
    int cur_len = 0;
    for (int i = 0; i < 8; ++i) {
        if (groups[i] == 0) {
            if (cur_start < 0) cur_start = i;
            ++cur_len;
            if (cur_len > best_len) {
                best_start = cur_start;
                best_len = cur_len;
            }
        } else {
            cur_start = -1;
            cur_len = 0;
        }
    }
    if (best_len < 2) best_start = -1;  // don't compress a lone zero group

    std::string out;
    char buf[6];
    for (int i = 0; i < 8; ++i) {
        if (i == best_start) {
            out += "::";
            i += best_len - 1;  // the for-loop's ++i advances past the run
            continue;
        }
        if (!out.empty() && out.back() != ':') out += ':';
        std::snprintf(buf, sizeof(buf), "%x", groups[i]);
        out += buf;
    }
    return out;
}

}  // namespace wireframe::net