// SPDX-License-Identifier: MPL-2.0
#pragma once
#include "vibeqc/tensor/distributed.hpp"

namespace vibeqc::tensor::distributed {
// Version 1 is a bounded, canonical manifest envelope. The receiver already
// has an authenticated Operation: compare every encoded field, not a hash or
// a native struct. CTF global indexing (doi:10.1016/j.jpdc.2014.06.002, sec.4)
// and TAMM dependent domains (doi:10.1063/5.0142433, III) motivate these fields.
struct WireHeader {
    static constexpr std::size_t capacity=4096;
    std::array<std::uint8_t,capacity> bytes{};
    std::size_t size=0;
    void integer(std::uint64_t value,unsigned width) {
        require(width<=8 && size<=capacity-width,"tensor wire header exceeds bound");
        if(width<8)require(value<(std::uint64_t(1)<<(8*width)),"wire integer narrowing");
        for(unsigned i=0;i<width;++i)bytes[size++]=std::uint8_t(value>>(8*(width-1-i)));
    }
    void id(const Id& v) { integer(std::uint16_t(v.kind),2);integer(v.namespace_high,8);integer(v.namespace_low,8);integer(v.serial,8); }
    void token(const Token& v) { id(v.id);integer(v.revision,8); }
    void domain(const DomainRef& v) {
        v.validate();token(v.token);integer(v.dependent,1);
        if(v.dependent) {token(v.parent);integer(v.parent_coordinate,8);}
        integer(v.selected,1);integer(v.extent,8);
        if(v.selected)for(std::size_t i=0;i<v.extent;++i)integer(v.coordinates[i],8);
        else {integer(v.begin,8);integer(v.step,8);}
    }
    void descriptor(const Descriptor& d) {
        d.validate();id(d.id);integer(d.schema,2);integer(d.structure_revision,8);
        integer(d.dtype==DType::fp64?1:2,2);integer(d.rank,2);
        token(d.source);token(d.metric);token(d.frame);
        for(std::size_t i=0;i<d.rank;++i)domain(d.axes[i]);
    }
};
inline WireHeader encode_header(const Operation& op,std::uint64_t p,std::uint64_t generation,std::uint32_t crc) {
    op.placement.validate();
    WireHeader h;
    h.integer(0x56515443,4); // VQTC
    h.integer(1,2);h.integer(0,2); // version, patched header length
    h.token(op.id);h.integer(op.value_epoch,8);
    h.descriptor(op.left);h.descriptor(op.right);h.descriptor(op.output);
    h.integer(op.placement.ranks,4);h.integer(op.placement.active,4);
    h.integer(op.placement.executor_shift,4);h.integer(op.placement.output_owner,4);h.integer(op.placement.epoch,8);
    h.integer(op.work_limit,8);h.integer(op.adjoint,1);
    // Version 1 fixes alpha=1, beta=0, compensated ordered CPU arithmetic,
    // full output block, and one contiguous p row. Unsupported modes refuse.
    h.integer(p,8);h.integer(generation,8);
    h.integer(op.placement.executor(p),4);h.integer(op.placement.output_owner,4);
    h.integer(mul(mul(op.output.axes[0].extent,op.output.axes[1].extent),16),8);
    h.integer(crc,4);
    h.bytes[6]=std::uint8_t(h.size>>8);h.bytes[7]=std::uint8_t(h.size);
    return h;
}
inline void validate_header(const WireHeader& expected,const std::uint8_t* data,std::size_t size) {
    require(data && size==expected.size && size>=8,"tensor header size mismatch");
    const auto declared=(std::size_t(data[6])<<8)|data[7];
    require(declared==size && declared<=WireHeader::capacity,"invalid bounded header length");
    require(std::equal(expected.bytes.begin(),expected.bytes.begin()+expected.size,data),"tensor envelope identity/schema/epoch mismatch");
}
inline void validate_packet(const Operation& op,std::uint64_t p,std::uint64_t generation,
                            const std::uint8_t* data,std::size_t size) {
    const auto initial=encode_header(op,p,generation,0);
    const auto payload=mul(mul(op.output.axes[0].extent,op.output.axes[1].extent),16);
    require(data && size==add(initial.size,payload),"tensor packet size mismatch");
    const auto expected=encode_header(op,p,generation,crc32(data+initial.size,payload));
    validate_header(expected,data,initial.size);
}
} // namespace vibeqc::tensor::distributed