#include "icmp.h"
#include <cstring>
#include "ipv4.h"

namespace icmp {

void handle(std::span<const uint8_t> frame, span_writer& tx,
            eth::mac_addr our_mac, ipv4::ip4_addr our_ip,
            std::function<void(std::span<const uint8_t>)> send_raw) {
    auto& ip = *reinterpret_cast<const ipv4::header*>(frame.data());
    size_t ip_hdr_len = ip.ip_header_len();
    size_t ip_total = ip.ip_total_len();

    if (sizeof(eth::header) + ip_total > frame.size()) return;
    if (ip.protocol != 1) return;

    auto& icmp_pkt = *reinterpret_cast<const echo*>(frame.data() + sizeof(eth::header) + ip_hdr_len);
    size_t icmp_len = ip_total - ip_hdr_len;
    if (icmp_len < sizeof(echo)) return;
    if (icmp_pkt.type != 8) return;

    size_t reply_len = sizeof(eth::header) + ip_total;
    if (reply_len > tx.capacity()) return;

    memcpy(tx.data(), frame.data(), reply_len);
    auto& rip = *reinterpret_cast<ipv4::header*>(tx.data());
    rip.eth.dst = ip.eth.src;
    rip.eth.src = our_mac;
    rip.src = our_ip;
    rip.dst = ip.src;
    rip.ttl = 64;
    rip.checksum = 0;
    rip.checksum = ipv4::checksum(rip.ip_start(), ip_hdr_len);

    auto& ricmp = *reinterpret_cast<echo*>(tx.data() + sizeof(eth::header) + ip_hdr_len);
    ricmp.type = 0;
    ricmp.checksum = 0;
    ricmp.checksum = ipv4::checksum(&ricmp, icmp_len);

    send_raw({tx.data(), reply_len});
}

size_t build_echo_request(std::span<uint8_t> buf,
                          eth::mac_addr src_mac, ipv4::ip4_addr src_ip,
                          eth::mac_addr dst_mac, ipv4::ip4_addr dst_ip,
                          uint16_t id, uint16_t seq) {
    size_t total = sizeof(ipv4::header) + sizeof(echo);
    if (buf.size() < total) return 0;

    memset(buf.data(), 0, total);

    auto& ip = *reinterpret_cast<ipv4::header*>(buf.data());
    ip.eth.dst = dst_mac;
    ip.eth.src = src_mac;
    ip.eth.ethertype = eth::ETH_IPV4;
    ip.ver_ihl = 0x45;
    ip.dscp_ecn = 0;
    ip.total_len = __builtin_bswap16(20 + sizeof(echo));
    ip.identification = 0;
    ip.flags_frag = 0;
    ip.ttl = 64;
    ip.protocol = 1;
    ip.checksum = 0;
    ip.src = src_ip;
    ip.dst = dst_ip;
    ip.checksum = ipv4::checksum(ip.ip_start(), 20);

    auto& icmp_pkt = *reinterpret_cast<echo*>(buf.data() + sizeof(ipv4::header));
    icmp_pkt.type = 8;
    icmp_pkt.code = 0;
    icmp_pkt.checksum = 0;
    icmp_pkt.id = id;
    icmp_pkt.seq = seq;
    icmp_pkt.checksum = ipv4::checksum(&icmp_pkt, sizeof(echo));

    return total;
}

bool parse_echo_reply(std::span<const uint8_t> frame, ipv4::ip4_addr& src_ip, uint16_t expected_id) {
    if (frame.size() < sizeof(ipv4::header)) return false;
    auto& ip = *reinterpret_cast<const ipv4::header*>(frame.data());
    if ((ip.ver_ihl >> 4) != 4) return false;
    if (ip.eth.ethertype != eth::ETH_IPV4) return false;
    if (ip.protocol != 1) return false;

    size_t ip_hdr_len = ip.ip_header_len();
    if (sizeof(eth::header) + ip_hdr_len + sizeof(echo) > frame.size()) return false;

    auto& icmp_pkt = *reinterpret_cast<const echo*>(frame.data() + sizeof(eth::header) + ip_hdr_len);
    if (icmp_pkt.type != 0) return false;
    if (icmp_pkt.id != expected_id) return false;

    src_ip = ip.src;
    return true;
}

} // namespace icmp