// STEP 9 deterministic native dataset pipeline acceptance executable.
//
// Input:
//   canonical P9CAN001 conversation binary created by normalize_dataset.py.
//
// Output:
//   deterministic response-only token masks, hash-stable train/validation
//   splits, fixed-length packed blocks, reset flags, causal-label validation,
//   checksum-protected P9DS0001 shards, and a JSON manifest.
//
// The model is loaded vocab-only. No base matrix is allocated or modified.

#include "llama.h"

#include <algorithm>
#include <array>
#include <chrono>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <set>
#include <sstream>
#include <stdexcept>
#include <string>
#include <thread>
#include <utility>
#include <vector>

namespace fs = std::filesystem;

struct message_record {
    std::string role;
    std::string content;
};

struct conversation_record {
    std::string id;
    std::string source_format;
    std::vector<message_record> messages;
};

struct tokenized_record {
    std::string id;
    std::string source_format;
    std::vector<llama_token> tokens;
    std::vector<uint8_t> target_mask;
    std::vector<uint8_t> expected_mask;
    uint64_t order_hash = 0;
    bool validation = false;
    bool truncated = false;
    uint32_t truncated_left = 0;
    uint32_t assistant_messages = 0;
    uint32_t supervised_tokens = 0;
    bool token_roundtrip = false;
    bool assistant_text_found = false;
};

struct packed_block {
    std::vector<int32_t> tokens;
    std::vector<uint8_t> target_mask;
    std::vector<uint8_t> sequence_start;
    std::vector<int32_t> source_index;
};

struct shard_info {
    fs::path path;
    std::string split_name;
    uint32_t block_count = 0;
    uint64_t payload_bytes = 0;
    uint64_t payload_fnv64 = 0;
    std::string sha256;
};

static const std::array<char, 8> CANONICAL_MAGIC = {'P','9','C','A','N','0','0','1'};
static const std::array<char, 8> SHARD_MAGIC     = {'P','9','D','S','0','0','0','1'};

static std::string shell_quote(const std::string & input) {
    std::string output = "'";
    for (const char value : input) {
        if (value == '\'') {
            output += "'\"'\"'";
        } else {
            output += value;
        }
    }
    output += "'";
    return output;
}

static std::string sha256_file(const fs::path & path) {
    const std::string command = "sha256sum " + shell_quote(path.string());
    FILE * pipe = popen(command.c_str(), "r");
    if (!pipe) {
        throw std::runtime_error("could not run sha256sum");
    }
    std::string output;
    char buffer[512] = {};
    while (fgets(buffer, sizeof(buffer), pipe)) {
        output += buffer;
    }
    const int rc = pclose(pipe);
    if (rc != 0 || output.size() < 64) {
        throw std::runtime_error("sha256sum failed for " + path.string());
    }
    return output.substr(0, 64);
}

static std::string json_escape(const std::string & input) {
    std::ostringstream output;
    for (const unsigned char value : input) {
        switch (value) {
            case '"':  output << "\\\""; break;
            case '\\': output << "\\\\"; break;
            case '\b': output << "\\b"; break;
            case '\f': output << "\\f"; break;
            case '\n': output << "\\n"; break;
            case '\r': output << "\\r"; break;
            case '\t': output << "\\t"; break;
            default:
                if (value < 0x20) {
                    output << "\\u"
                           << std::hex << std::setw(4) << std::setfill('0')
                           << static_cast<int>(value)
                           << std::dec << std::setfill(' ');
                } else {
                    output << static_cast<char>(value);
                }
        }
    }
    return output.str();
}

static uint64_t fnv1a64_bytes(
        const uint8_t * data,
        size_t size,
        uint64_t seed = 1469598103934665603ULL) {
    uint64_t value = seed;
    for (size_t index = 0; index < size; ++index) {
        value ^= static_cast<uint64_t>(data[index]);
        value *= 1099511628211ULL;
    }
    return value;
}

static uint64_t stable_record_hash(const std::string & id, uint64_t seed) {
    uint64_t value = 1469598103934665603ULL;
    for (int index = 0; index < 8; ++index) {
        const uint8_t byte = static_cast<uint8_t>((seed >> (8 * index)) & 0xff);
        value ^= byte;
        value *= 1099511628211ULL;
    }
    return fnv1a64_bytes(
        reinterpret_cast<const uint8_t *>(id.data()),
        id.size(),
        value);
}

static int gpu_memory_mib() {
    FILE * pipe = popen(
        "nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null",
        "r");
    if (!pipe) {
        return -1;
    }
    char buffer[128] = {};
    int total = 0;
    bool any = false;
    while (fgets(buffer, sizeof(buffer), pipe)) {
        total += std::atoi(buffer);
        any = true;
    }
    pclose(pipe);
    return any ? total : -1;
}

template<typename T>
static void append_scalar(std::vector<uint8_t> & output, T value) {
    for (size_t index = 0; index < sizeof(T); ++index) {
        output.push_back(static_cast<uint8_t>(
            (static_cast<uint64_t>(value) >> (8 * index)) & 0xff));
    }
}

static uint32_t read_u32(std::istream & input) {
    uint8_t bytes[4] = {};
    input.read(reinterpret_cast<char *>(bytes), 4);
    if (!input) {
        throw std::runtime_error("unexpected EOF reading u32");
    }
    return
        static_cast<uint32_t>(bytes[0]) |
        (static_cast<uint32_t>(bytes[1]) << 8) |
        (static_cast<uint32_t>(bytes[2]) << 16) |
        (static_cast<uint32_t>(bytes[3]) << 24);
}

static uint64_t read_u64(std::istream & input) {
    uint8_t bytes[8] = {};
    input.read(reinterpret_cast<char *>(bytes), 8);
    if (!input) {
        throw std::runtime_error("unexpected EOF reading u64");
    }
    uint64_t value = 0;
    for (int index = 0; index < 8; ++index) {
        value |= static_cast<uint64_t>(bytes[index]) << (8 * index);
    }
    return value;
}

static int32_t read_i32(std::istream & input) {
    return static_cast<int32_t>(read_u32(input));
}

static std::string read_string(std::istream & input) {
    const uint32_t size = read_u32(input);
    if (size > 256U * 1024U * 1024U) {
        throw std::runtime_error("unreasonable string length");
    }
    std::string value(size, '\0');
    input.read(value.data(), static_cast<std::streamsize>(size));
    if (!input) {
        throw std::runtime_error("unexpected EOF reading string");
    }
    return value;
}

static void write_u32(std::ostream & output, uint32_t value) {
    uint8_t bytes[4] = {
        static_cast<uint8_t>(value & 0xff),
        static_cast<uint8_t>((value >> 8) & 0xff),
        static_cast<uint8_t>((value >> 16) & 0xff),
        static_cast<uint8_t>((value >> 24) & 0xff),
    };
    output.write(reinterpret_cast<const char *>(bytes), 4);
}

static void write_u64(std::ostream & output, uint64_t value) {
    uint8_t bytes[8] = {};
    for (int index = 0; index < 8; ++index) {
        bytes[index] = static_cast<uint8_t>((value >> (8 * index)) & 0xff);
    }
    output.write(reinterpret_cast<const char *>(bytes), 8);
}

static void write_i32(std::ostream & output, int32_t value) {
    write_u32(output, static_cast<uint32_t>(value));
}

static std::vector<conversation_record> read_canonical(const fs::path & path) {
    std::ifstream input(path, std::ios::binary);
    if (!input) {
        throw std::runtime_error("could not open canonical dataset: " + path.string());
    }

    std::array<char, 8> magic = {};
    input.read(magic.data(), static_cast<std::streamsize>(magic.size()));
    if (!input || magic != CANONICAL_MAGIC) {
        throw std::runtime_error("invalid canonical dataset magic");
    }

    const uint32_t version = read_u32(input);
    if (version != 1) {
        throw std::runtime_error("unsupported canonical dataset version");
    }

    const uint32_t record_count = read_u32(input);
    if (record_count == 0 || record_count > 10000000U) {
        throw std::runtime_error("invalid canonical record count");
    }

    std::vector<conversation_record> records;
    records.reserve(record_count);
    std::set<std::string> ids;

    for (uint32_t record_index = 0; record_index < record_count; ++record_index) {
        conversation_record record;
        record.id = read_string(input);
        record.source_format = read_string(input);
        const uint32_t message_count = read_u32(input);

        if (record.id.empty() || message_count == 0 || message_count > 100000U) {
            throw std::runtime_error("invalid canonical record");
        }
        if (!ids.insert(record.id).second) {
            throw std::runtime_error("duplicate canonical record id: " + record.id);
        }

        uint32_t assistant_count = 0;
        for (uint32_t message_index = 0; message_index < message_count; ++message_index) {
            message_record message;
            message.role = read_string(input);
            message.content = read_string(input);
            if (
                message.role != "system" &&
                message.role != "user" &&
                message.role != "assistant") {
                throw std::runtime_error("unsupported message role in canonical input");
            }
            if (message.role == "assistant") {
                if (message.content.empty()) {
                    throw std::runtime_error("empty assistant content");
                }
                ++assistant_count;
            }
            record.messages.push_back(std::move(message));
        }

        if (assistant_count == 0) {
            throw std::runtime_error("record has no assistant message: " + record.id);
        }
        records.push_back(std::move(record));
    }

    char trailing = 0;
    if (input.read(&trailing, 1)) {
        throw std::runtime_error("canonical dataset has trailing bytes");
    }

    return records;
}

static std::string apply_template(
        const char * chat_template,
        const std::vector<message_record> & messages,
        bool add_assistant) {
    std::vector<llama_chat_message> chat;
    chat.reserve(messages.size());
    for (const auto & message : messages) {
        chat.push_back({message.role.c_str(), message.content.c_str()});
    }

    size_t estimated = 1024;
    for (const auto & message : messages) {
        estimated += message.role.size() + message.content.size() + 64;
    }

    std::vector<char> buffer(estimated);
    int32_t result = llama_chat_apply_template(
        chat_template,
        chat.data(),
        chat.size(),
        add_assistant,
        buffer.data(),
        static_cast<int32_t>(buffer.size()));

    if (result < 0) {
        throw std::runtime_error("llama_chat_apply_template returned an error");
    }

    if (static_cast<size_t>(result) >= buffer.size()) {
        buffer.resize(static_cast<size_t>(result) + 1);
        result = llama_chat_apply_template(
            chat_template,
            chat.data(),
            chat.size(),
            add_assistant,
            buffer.data(),
            static_cast<int32_t>(buffer.size()));
        if (result < 0 || static_cast<size_t>(result) >= buffer.size()) {
            throw std::runtime_error("chat template reallocation failed");
        }
    }

    return std::string(buffer.data(), static_cast<size_t>(result));
}

static std::vector<llama_token> tokenize_text(
        const llama_vocab * vocab,
        const std::string & text) {
    int32_t capacity = std::max<int32_t>(
        32,
        static_cast<int32_t>(text.size() + 16));
    std::vector<llama_token> tokens(static_cast<size_t>(capacity));

    int32_t count = llama_tokenize(
        vocab,
        text.data(),
        static_cast<int32_t>(text.size()),
        tokens.data(),
        capacity,
        false,
        true);

    if (count == std::numeric_limits<int32_t>::min()) {
        throw std::runtime_error("tokenization overflow");
    }

    if (count < 0) {
        capacity = -count;
        tokens.resize(static_cast<size_t>(capacity));
        count = llama_tokenize(
            vocab,
            text.data(),
            static_cast<int32_t>(text.size()),
            tokens.data(),
            capacity,
            false,
            true);
    }

    if (count < 0) {
        throw std::runtime_error("tokenization failed");
    }

    tokens.resize(static_cast<size_t>(count));
    return tokens;
}

static std::string detokenize_text(
        const llama_vocab * vocab,
        const std::vector<llama_token> & tokens) {
    int32_t capacity = std::max<int32_t>(
        64,
        static_cast<int32_t>(tokens.size() * 16 + 64));
    std::vector<char> buffer(static_cast<size_t>(capacity));

    int32_t count = llama_detokenize(
        vocab,
        tokens.data(),
        static_cast<int32_t>(tokens.size()),
        buffer.data(),
        capacity,
        false,
        true);

    if (count < 0) {
        capacity = -count;
        buffer.resize(static_cast<size_t>(capacity));
        count = llama_detokenize(
            vocab,
            tokens.data(),
            static_cast<int32_t>(tokens.size()),
            buffer.data(),
            capacity,
            false,
            true);
    }

    if (count < 0) {
        throw std::runtime_error("detokenization failed");
    }

    return std::string(buffer.data(), static_cast<size_t>(count));
}

static bool is_token_prefix(
        const std::vector<llama_token> & prefix,
        const std::vector<llama_token> & full) {
    return
        prefix.size() <= full.size() &&
        std::equal(prefix.begin(), prefix.end(), full.begin());
}

static uint32_t count_mask(const std::vector<uint8_t> & mask) {
    uint32_t count = 0;
    for (const uint8_t value : mask) {
        count += value != 0;
    }
    return count;
}

static tokenized_record tokenize_record(
        const conversation_record & input,
        const llama_vocab * vocab,
        const char * chat_template,
        int32_t vocab_size,
        uint32_t max_sequence_length,
        uint64_t seed,
        uint32_t validation_permille) {
    tokenized_record output;
    output.id = input.id;
    output.source_format = input.source_format;
    output.order_hash = stable_record_hash(input.id, seed);
    output.validation = output.order_hash % 1000ULL < validation_permille;

    const std::string final_text =
        apply_template(chat_template, input.messages, false);
    output.tokens = tokenize_text(vocab, final_text);
    output.target_mask.assign(output.tokens.size(), 0);
    output.expected_mask.assign(output.tokens.size(), 0);

    if (output.tokens.empty()) {
        throw std::runtime_error("record tokenized to zero tokens: " + input.id);
    }

    for (const llama_token token : output.tokens) {
        if (token < 0 || token >= vocab_size) {
            throw std::runtime_error("token id out of vocabulary range");
        }
    }

    std::vector<std::pair<size_t, size_t>> assistant_spans;

    for (size_t index = 0; index < input.messages.size(); ++index) {
        if (input.messages[index].role != "assistant") {
            continue;
        }

        ++output.assistant_messages;

        std::vector<message_record> before(
            input.messages.begin(),
            input.messages.begin() + static_cast<std::ptrdiff_t>(index));
        std::vector<message_record> through(
            input.messages.begin(),
            input.messages.begin() + static_cast<std::ptrdiff_t>(index + 1));

        const std::string prefix_text =
            apply_template(chat_template, before, true);
        const std::string through_text =
            apply_template(chat_template, through, false);

        const std::vector<llama_token> prefix_tokens =
            tokenize_text(vocab, prefix_text);
        const std::vector<llama_token> through_tokens =
            tokenize_text(vocab, through_text);

        if (!is_token_prefix(prefix_tokens, through_tokens)) {
            throw std::runtime_error(
                "assistant prefix is not token-prefix-stable for " + input.id);
        }
        if (!is_token_prefix(through_tokens, output.tokens)) {
            throw std::runtime_error(
                "conversation prefix is not token-prefix-stable for " + input.id);
        }
        if (prefix_tokens.size() >= through_tokens.size()) {
            throw std::runtime_error(
                "assistant span contains no tokens for " + input.id);
        }

        assistant_spans.push_back({
            prefix_tokens.size(),
            through_tokens.size(),
        });

        for (size_t token_index = prefix_tokens.size();
             token_index < through_tokens.size();
             ++token_index) {
            output.target_mask[token_index] = 1;
            output.expected_mask[token_index] = 1;
        }

        std::vector<llama_token> span_tokens(
            through_tokens.begin() + static_cast<std::ptrdiff_t>(prefix_tokens.size()),
            through_tokens.end());
        const std::string span_text = detokenize_text(vocab, span_tokens);
        if (span_text.find(input.messages[index].content) != std::string::npos) {
            output.assistant_text_found = true;
        }
    }

    if (output.assistant_messages == 0) {
        throw std::runtime_error("no assistant messages after native parsing");
    }

    const std::string detokenized = detokenize_text(vocab, output.tokens);
    const std::vector<llama_token> retokenized =
        tokenize_text(vocab, detokenized);
    output.token_roundtrip = retokenized == output.tokens;
    if (!output.token_roundtrip) {
        throw std::runtime_error("token/detokenize roundtrip mismatch");
    }

    if (output.tokens.size() > max_sequence_length) {
        const size_t remove_count =
            output.tokens.size() - max_sequence_length;
        output.tokens.erase(
            output.tokens.begin(),
            output.tokens.begin() + static_cast<std::ptrdiff_t>(remove_count));
        output.target_mask.erase(
            output.target_mask.begin(),
            output.target_mask.begin() + static_cast<std::ptrdiff_t>(remove_count));
        output.expected_mask.erase(
            output.expected_mask.begin(),
            output.expected_mask.begin() + static_cast<std::ptrdiff_t>(remove_count));
        output.truncated = true;
        output.truncated_left = static_cast<uint32_t>(remove_count);
    }

    // A token at position zero has no preceding logit in this isolated sample.
    if (!output.target_mask.empty()) {
        output.target_mask[0] = 0;
        output.expected_mask[0] = 0;
    }

    output.supervised_tokens = count_mask(output.target_mask);
    if (output.supervised_tokens == 0) {
        throw std::runtime_error(
            "truncation removed all supervised tokens for " + input.id);
    }

    if (output.target_mask != output.expected_mask) {
        throw std::runtime_error("response-only mask mismatch");
    }

    return output;
}

static std::vector<packed_block> pack_records(
        const std::vector<tokenized_record> & records,
        bool validation,
        uint32_t max_sequence_length,
        int32_t pad_token) {
    std::vector<size_t> selected;
    for (size_t index = 0; index < records.size(); ++index) {
        if (records[index].validation == validation) {
            selected.push_back(index);
        }
    }

    std::sort(
        selected.begin(),
        selected.end(),
        [&](size_t left, size_t right) {
            if (records[left].order_hash != records[right].order_hash) {
                return records[left].order_hash < records[right].order_hash;
            }
            return records[left].id < records[right].id;
        });

    std::vector<packed_block> blocks;
    packed_block current;

    auto initialize = [&]() {
        current.tokens.clear();
        current.target_mask.clear();
        current.sequence_start.clear();
        current.source_index.clear();
    };

    auto flush = [&]() {
        if (current.tokens.empty()) {
            return;
        }
        while (current.tokens.size() < max_sequence_length) {
            current.tokens.push_back(pad_token);
            current.target_mask.push_back(0);
            current.sequence_start.push_back(0);
            current.source_index.push_back(-1);
        }
        blocks.push_back(current);
        initialize();
    };

    initialize();

    for (const size_t record_index : selected) {
        const tokenized_record & record = records[record_index];
        if (record.tokens.size() > max_sequence_length) {
            throw std::runtime_error("record exceeds maximum after truncation");
        }
        if (
            !current.tokens.empty() &&
            current.tokens.size() + record.tokens.size() > max_sequence_length) {
            flush();
        }

        const size_t offset = current.tokens.size();
        for (size_t index = 0; index < record.tokens.size(); ++index) {
            current.tokens.push_back(record.tokens[index]);
            current.target_mask.push_back(record.target_mask[index]);
            current.sequence_start.push_back(index == 0 ? 1 : 0);
            current.source_index.push_back(static_cast<int32_t>(record_index));
        }

        if (offset >= current.sequence_start.size() ||
            current.sequence_start[offset] != 1) {
            throw std::runtime_error("sample boundary reset flag missing");
        }
        if (current.target_mask[offset] != 0) {
            throw std::runtime_error("sample-first token must not be supervised");
        }
    }

    flush();
    return blocks;
}

static std::vector<uint8_t> serialize_blocks(
        const std::vector<packed_block> & blocks) {
    std::vector<uint8_t> payload;

    for (size_t block_index = 0; block_index < blocks.size(); ++block_index) {
        const packed_block & block = blocks[block_index];
        append_scalar<uint32_t>(payload, static_cast<uint32_t>(block_index));
        append_scalar<uint32_t>(payload, static_cast<uint32_t>(block.tokens.size()));

        for (const int32_t value : block.tokens) {
            append_scalar<uint32_t>(payload, static_cast<uint32_t>(value));
        }
        payload.insert(
            payload.end(),
            block.target_mask.begin(),
            block.target_mask.end());
        payload.insert(
            payload.end(),
            block.sequence_start.begin(),
            block.sequence_start.end());
        for (const int32_t value : block.source_index) {
            append_scalar<uint32_t>(payload, static_cast<uint32_t>(value));
        }
    }

    return payload;
}

static shard_info write_shard(
        const fs::path & path,
        const std::string & split_name,
        uint32_t split_id,
        const std::vector<packed_block> & blocks,
        uint32_t max_sequence_length,
        int32_t pad_token,
        uint64_t seed) {
    const std::vector<uint8_t> payload = serialize_blocks(blocks);
    const uint64_t checksum =
        fnv1a64_bytes(payload.data(), payload.size());

    fs::create_directories(path.parent_path());
    std::ofstream output(path, std::ios::binary);
    if (!output) {
        throw std::runtime_error("could not create shard");
    }

    output.write(SHARD_MAGIC.data(), SHARD_MAGIC.size());
    write_u32(output, 1);
    write_u32(output, split_id);
    write_u32(output, max_sequence_length);
    write_u32(output, static_cast<uint32_t>(blocks.size()));
    write_i32(output, pad_token);
    write_u64(output, seed);
    write_u64(output, static_cast<uint64_t>(payload.size()));
    write_u64(output, checksum);
    output.write(
        reinterpret_cast<const char *>(payload.data()),
        static_cast<std::streamsize>(payload.size()));

    if (!output) {
        throw std::runtime_error("failed writing shard");
    }
    output.close();

    shard_info info;
    info.path = path;
    info.split_name = split_name;
    info.block_count = static_cast<uint32_t>(blocks.size());
    info.payload_bytes = payload.size();
    info.payload_fnv64 = checksum;
    info.sha256 = sha256_file(path);
    return info;
}

static std::vector<packed_block> read_shard(
        const fs::path & path,
        uint32_t expected_max_sequence_length,
        uint64_t expected_seed,
        bool * checksum_mismatch) {
    if (checksum_mismatch) {
        *checksum_mismatch = false;
    }

    std::ifstream input(path, std::ios::binary);
    if (!input) {
        throw std::runtime_error("could not open shard: " + path.string());
    }

    std::array<char, 8> magic = {};
    input.read(magic.data(), magic.size());
    if (!input || magic != SHARD_MAGIC) {
        throw std::runtime_error("invalid shard magic");
    }

    const uint32_t version = read_u32(input);
    const uint32_t split_id = read_u32(input);
    const uint32_t max_sequence_length = read_u32(input);
    const uint32_t block_count = read_u32(input);
    const int32_t pad_token = read_i32(input);
    const uint64_t seed = read_u64(input);
    const uint64_t payload_size = read_u64(input);
    const uint64_t stored_checksum = read_u64(input);

    (void) split_id;
    (void) pad_token;

    if (version != 1) {
        throw std::runtime_error("unsupported shard version");
    }
    if (max_sequence_length != expected_max_sequence_length) {
        throw std::runtime_error("shard maximum length mismatch");
    }
    if (seed != expected_seed) {
        throw std::runtime_error("shard seed mismatch");
    }
    if (payload_size > 16ULL * 1024ULL * 1024ULL * 1024ULL) {
        throw std::runtime_error("unreasonable shard payload size");
    }

    std::vector<uint8_t> payload(static_cast<size_t>(payload_size));
    input.read(
        reinterpret_cast<char *>(payload.data()),
        static_cast<std::streamsize>(payload.size()));
    if (!input) {
        throw std::runtime_error("truncated shard payload");
    }
    char trailing = 0;
    if (input.read(&trailing, 1)) {
        throw std::runtime_error("shard has trailing bytes");
    }

    const uint64_t actual_checksum =
        fnv1a64_bytes(payload.data(), payload.size());
    if (actual_checksum != stored_checksum) {
        if (checksum_mismatch) {
            *checksum_mismatch = true;
        }
        throw std::runtime_error("shard checksum mismatch");
    }

    size_t offset = 0;

    auto take_u32 = [&]() -> uint32_t {
        if (offset + 4 > payload.size()) {
            throw std::runtime_error("payload EOF reading u32");
        }
        const uint32_t value =
            static_cast<uint32_t>(payload[offset]) |
            (static_cast<uint32_t>(payload[offset + 1]) << 8) |
            (static_cast<uint32_t>(payload[offset + 2]) << 16) |
            (static_cast<uint32_t>(payload[offset + 3]) << 24);
        offset += 4;
        return value;
    };

    std::vector<packed_block> blocks;
    blocks.reserve(block_count);

    for (uint32_t expected_block = 0; expected_block < block_count; ++expected_block) {
        const uint32_t block_index = take_u32();
        const uint32_t token_count = take_u32();
        if (block_index != expected_block) {
            throw std::runtime_error("shard block index mismatch");
        }
        if (token_count != max_sequence_length) {
            throw std::runtime_error("packed block length mismatch");
        }

        packed_block block;
        block.tokens.reserve(token_count);
        block.target_mask.resize(token_count);
        block.sequence_start.resize(token_count);
        block.source_index.reserve(token_count);

        for (uint32_t index = 0; index < token_count; ++index) {
            block.tokens.push_back(static_cast<int32_t>(take_u32()));
        }

        if (offset + token_count > payload.size()) {
            throw std::runtime_error("payload EOF reading target mask");
        }
        std::copy(
            payload.begin() + static_cast<std::ptrdiff_t>(offset),
            payload.begin() + static_cast<std::ptrdiff_t>(offset + token_count),
            block.target_mask.begin());
        offset += token_count;

        if (offset + token_count > payload.size()) {
            throw std::runtime_error("payload EOF reading reset mask");
        }
        std::copy(
            payload.begin() + static_cast<std::ptrdiff_t>(offset),
            payload.begin() + static_cast<std::ptrdiff_t>(offset + token_count),
            block.sequence_start.begin());
        offset += token_count;

        for (uint32_t index = 0; index < token_count; ++index) {
            block.source_index.push_back(static_cast<int32_t>(take_u32()));
        }

        blocks.push_back(std::move(block));
    }

    if (offset != payload.size()) {
        throw std::runtime_error("unconsumed shard payload bytes");
    }

    return blocks;
}

static bool blocks_equal(
        const std::vector<packed_block> & left,
        const std::vector<packed_block> & right) {
    if (left.size() != right.size()) {
        return false;
    }
    for (size_t index = 0; index < left.size(); ++index) {
        if (
            left[index].tokens != right[index].tokens ||
            left[index].target_mask != right[index].target_mask ||
            left[index].sequence_start != right[index].sequence_start ||
            left[index].source_index != right[index].source_index) {
            return false;
        }
    }
    return true;
}

struct pipeline_result {
    std::vector<tokenized_record> records;
    std::vector<packed_block> train_blocks;
    std::vector<packed_block> validation_blocks;
    std::vector<shard_info> shards;
    uint32_t train_records = 0;
    uint32_t validation_records = 0;
    uint32_t assistant_messages = 0;
    uint32_t supervised_tokens = 0;
    uint32_t truncated_records = 0;
    uint32_t prompt_mask_violations = 0;
    uint32_t padding_mask_violations = 0;
    uint32_t boundary_violations = 0;
    uint32_t label_violations = 0;
    uint32_t token_range_violations = 0;
    uint32_t roundtrip_violations = 0;
    uint32_t assistant_text_violations = 0;
    bool shard_reload_exact = false;
    std::string fingerprint;
};

static std::vector<shard_info> write_split_shards(
        const fs::path & output_dir,
        const std::string & split_name,
        uint32_t split_id,
        const std::vector<packed_block> & blocks,
        uint32_t max_sequence_length,
        int32_t pad_token,
        uint64_t seed,
        uint32_t shard_block_capacity) {
    if (shard_block_capacity == 0) {
        throw std::runtime_error("shard block capacity must be positive");
    }

    std::vector<shard_info> result;
    for (size_t begin = 0, shard_index = 0;
         begin < blocks.size();
         begin += shard_block_capacity, ++shard_index) {
        const size_t end = std::min(
            blocks.size(),
            begin + static_cast<size_t>(shard_block_capacity));
        std::vector<packed_block> slice(
            blocks.begin() + static_cast<std::ptrdiff_t>(begin),
            blocks.begin() + static_cast<std::ptrdiff_t>(end));

        std::ostringstream filename;
        filename
            << split_name << '-'
            << std::setw(5) << std::setfill('0') << shard_index
            << ".p9ds";

        result.push_back(write_shard(
            output_dir / filename.str(),
            split_name,
            split_id,
            slice,
            max_sequence_length,
            pad_token,
            seed));
    }
    return result;
}

static std::vector<packed_block> reload_split_shards(
        const std::vector<shard_info> & shards,
        const std::string & split_name,
        uint32_t max_sequence_length,
        uint64_t seed,
        bool * checksum_mismatch) {
    if (checksum_mismatch) {
        *checksum_mismatch = false;
    }
    std::vector<packed_block> result;
    for (const shard_info & shard : shards) {
        if (shard.split_name != split_name) {
            continue;
        }
        bool local_mismatch = false;
        const std::vector<packed_block> loaded = read_shard(
            shard.path,
            max_sequence_length,
            seed,
            &local_mismatch);
        if (local_mismatch && checksum_mismatch) {
            *checksum_mismatch = true;
        }
        result.insert(result.end(), loaded.begin(), loaded.end());
    }
    return result;
}

static pipeline_result build_pipeline(
        const fs::path & model_path,
        const fs::path & canonical_path,
        const fs::path & output_dir,
        uint64_t seed,
        uint32_t max_sequence_length,
        uint32_t validation_permille,
        uint32_t shard_block_capacity,
        bool write_outputs) {

    fs::create_directories(output_dir);

    const int gpu_before = gpu_memory_mib();

    llama_backend_init();

    llama_model_params params = llama_model_default_params();
    params.vocab_only = true;
    params.n_gpu_layers = 0;
    params.use_mmap = true;
    params.check_tensors = false;

    llama_model * model =
        llama_model_load_from_file(model_path.c_str(), params);
    if (!model) {
        llama_backend_free();
        throw std::runtime_error("could not load model vocabulary");
    }

    const llama_vocab * vocab = llama_model_get_vocab(model);
    if (!vocab) {
        llama_model_free(model);
        llama_backend_free();
        throw std::runtime_error("model has no vocabulary");
    }

    const int32_t vocab_size = llama_vocab_n_tokens(vocab);
    const char * chat_template = llama_model_chat_template(model, nullptr);
    if (!chat_template || std::strlen(chat_template) == 0) {
        llama_model_free(model);
        llama_backend_free();
        throw std::runtime_error("model has no default chat template");
    }

    std::vector<conversation_record> conversations =
        read_canonical(canonical_path);
    std::sort(
        conversations.begin(),
        conversations.end(),
        [](const conversation_record & left, const conversation_record & right) {
            return left.id < right.id;
        });

    pipeline_result result;
    result.records.reserve(conversations.size());

    for (const conversation_record & conversation : conversations) {
        tokenized_record record = tokenize_record(
            conversation,
            vocab,
            chat_template,
            vocab_size,
            max_sequence_length,
            seed,
            validation_permille);

        result.assistant_messages += record.assistant_messages;
        result.supervised_tokens += record.supervised_tokens;
        result.truncated_records += record.truncated ? 1U : 0U;
        result.train_records += record.validation ? 0U : 1U;
        result.validation_records += record.validation ? 1U : 0U;

        if (record.target_mask != record.expected_mask) {
            ++result.prompt_mask_violations;
        }
        if (!record.token_roundtrip) {
            ++result.roundtrip_violations;
        }
        if (!record.assistant_text_found) {
            ++result.assistant_text_violations;
        }
        for (const llama_token token : record.tokens) {
            if (token < 0 || token >= vocab_size) {
                ++result.token_range_violations;
            }
        }

        result.records.push_back(std::move(record));
    }

    if (result.train_records == 0 || result.validation_records == 0) {
        llama_model_free(model);
        llama_backend_free();
        throw std::runtime_error("train or validation split is empty");
    }

    int32_t pad_token = llama_vocab_pad(vocab);
    if (pad_token < 0) {
        pad_token = llama_vocab_eos(vocab);
    }
    if (pad_token < 0) {
        pad_token = 0;
    }

    result.train_blocks = pack_records(
        result.records,
        false,
        max_sequence_length,
        pad_token);
    result.validation_blocks = pack_records(
        result.records,
        true,
        max_sequence_length,
        pad_token);

    auto validate_blocks = [&](const std::vector<packed_block> & blocks) {
        for (const packed_block & block : blocks) {
            if (
                block.tokens.size() != max_sequence_length ||
                block.target_mask.size() != max_sequence_length ||
                block.sequence_start.size() != max_sequence_length ||
                block.source_index.size() != max_sequence_length) {
                ++result.boundary_violations;
                continue;
            }

            for (size_t index = 0; index < block.tokens.size(); ++index) {
                const bool padding = block.source_index[index] < 0;
                if (padding && block.target_mask[index] != 0) {
                    ++result.padding_mask_violations;
                }
                if (
                    block.sequence_start[index] != 0 &&
                    block.target_mask[index] != 0) {
                    ++result.boundary_violations;
                }

                if (index + 1 < block.tokens.size()) {
                    const bool expected_label =
                        block.target_mask[index + 1] != 0 &&
                        block.sequence_start[index + 1] == 0 &&
                        block.source_index[index + 1] >= 0;
                    const int32_t label =
                        expected_label ? block.tokens[index + 1] : -100;
                    if (
                        expected_label && label < 0) {
                        ++result.label_violations;
                    }
                    if (
                        !expected_label && label != -100) {
                        ++result.label_violations;
                    }
                }
            }
        }
    };

    validate_blocks(result.train_blocks);
    validate_blocks(result.validation_blocks);

    if (write_outputs) {
        const std::vector<shard_info> train_shards = write_split_shards(
            output_dir,
            "train",
            0,
            result.train_blocks,
            max_sequence_length,
            pad_token,
            seed,
            shard_block_capacity);
        const std::vector<shard_info> validation_shards = write_split_shards(
            output_dir,
            "validation",
            1,
            result.validation_blocks,
            max_sequence_length,
            pad_token,
            seed,
            shard_block_capacity);

        result.shards.insert(
            result.shards.end(),
            train_shards.begin(),
            train_shards.end());
        result.shards.insert(
            result.shards.end(),
            validation_shards.begin(),
            validation_shards.end());

        bool train_checksum_mismatch = false;
        bool validation_checksum_mismatch = false;
        const std::vector<packed_block> train_reloaded = reload_split_shards(
            result.shards,
            "train",
            max_sequence_length,
            seed,
            &train_checksum_mismatch);
        const std::vector<packed_block> validation_reloaded = reload_split_shards(
            result.shards,
            "validation",
            max_sequence_length,
            seed,
            &validation_checksum_mismatch);

        result.shard_reload_exact =
            !train_checksum_mismatch &&
            !validation_checksum_mismatch &&
            blocks_equal(result.train_blocks, train_reloaded) &&
            blocks_equal(result.validation_blocks, validation_reloaded);

        const fs::path template_path = output_dir / "chat_template.txt";
        {
            std::ofstream output(template_path, std::ios::binary);
            output << chat_template;
        }

        std::vector<std::string> train_ids;
        std::vector<std::string> validation_ids;
        for (const tokenized_record & record : result.records) {
            (record.validation ? validation_ids : train_ids).push_back(record.id);
        }
        std::sort(train_ids.begin(), train_ids.end());
        std::sort(validation_ids.begin(), validation_ids.end());

        const fs::path split_path = output_dir / "split_assignments.tsv";
        {
            std::ofstream output(split_path);
            output << "id\tsplit\torder_hash\ttokens\tsupervised\ttruncated\n";
            std::vector<const tokenized_record *> ordered;
            for (const tokenized_record & record : result.records) {
                ordered.push_back(&record);
            }
            std::sort(
                ordered.begin(),
                ordered.end(),
                [](const tokenized_record * left, const tokenized_record * right) {
                    return left->id < right->id;
                });
            for (const tokenized_record * record : ordered) {
                output
                    << record->id << '\t'
                    << (record->validation ? "validation" : "train") << '\t'
                    << record->order_hash << '\t'
                    << record->tokens.size() << '\t'
                    << record->supervised_tokens << '\t'
                    << (record->truncated ? 1 : 0)
                    << '\n';
            }
        }

        const fs::path material_path = output_dir / "fingerprint_material.txt";
        {
            std::ofstream output(material_path);
            output << "format=prism-step9-dataset-v1\n";
            output << "seed=" << seed << "\n";
            output << "max_sequence_length=" << max_sequence_length << "\n";
            output << "validation_permille=" << validation_permille << "\n";
            output << "chat_template_sha256=" << sha256_file(template_path) << "\n";
            for (const shard_info & shard : result.shards) {
                output
                    << shard.split_name << '='
                    << shard.sha256 << ':'
                    << shard.payload_fnv64 << ':'
                    << shard.block_count << '\n';
            }
            for (const std::string & id : train_ids) {
                output << "train_id=" << id << "\n";
            }
            for (const std::string & id : validation_ids) {
                output << "validation_id=" << id << "\n";
            }
        }

        result.fingerprint = sha256_file(material_path);

        const fs::path manifest_path = output_dir / "dataset_manifest.json";
        std::ofstream manifest(manifest_path);
        manifest << "{\n";
        manifest << "  \"format\":\"prism-step9-dataset-v1\",\n";
        manifest << "  \"model\":\"" << json_escape(model_path.string()) << "\",\n";
        manifest << "  \"canonical\":\"" << json_escape(canonical_path.string()) << "\",\n";
        manifest << "  \"seed\":" << seed << ",\n";
        manifest << "  \"max_sequence_length\":" << max_sequence_length << ",\n";
        manifest << "  \"validation_permille\":" << validation_permille << ",\n";
        manifest << "  \"vocab_size\":" << vocab_size << ",\n";
        manifest << "  \"chat_template_sha256\":\""
                 << sha256_file(template_path) << "\",\n";
        manifest << "  \"record_count\":" << result.records.size() << ",\n";
        manifest << "  \"train_record_count\":" << result.train_records << ",\n";
        manifest << "  \"validation_record_count\":" << result.validation_records << ",\n";
        manifest << "  \"train_block_count\":" << result.train_blocks.size() << ",\n";
        manifest << "  \"validation_block_count\":" << result.validation_blocks.size() << ",\n";
        manifest << "  \"supervised_token_count\":" << result.supervised_tokens << ",\n";
        manifest << "  \"dataset_fingerprint\":\"" << result.fingerprint << "\",\n";
        manifest << "  \"shards\":[\n";
        for (size_t index = 0; index < result.shards.size(); ++index) {
            const shard_info & shard = result.shards[index];
            manifest
                << "    {\"split\":\"" << shard.split_name
                << "\",\"path\":\"" << json_escape(shard.path.string())
                << "\",\"sha256\":\"" << shard.sha256
                << "\",\"payload_fnv64\":" << shard.payload_fnv64
                << ",\"block_count\":" << shard.block_count
                << "}";
            if (index + 1 != result.shards.size()) {
                manifest << ',';
            }
            manifest << '\n';
        }
        manifest << "  ]\n";
        manifest << "}\n";
    }

    const std::string template_sha = write_outputs
        ? sha256_file(output_dir / "chat_template.txt")
        : std::string();

    llama_model_free(model);
    llama_backend_free();

    std::this_thread::sleep_for(std::chrono::milliseconds(250));
    const int gpu_after = gpu_memory_mib();

    std::cout << "MODEL_VOCAB_ONLY=1\n";
    std::cout << "VOCAB_SIZE=" << vocab_size << "\n";
    std::cout << "CHAT_TEMPLATE_PRESENT=1\n";
    if (!template_sha.empty()) {
        std::cout << "CHAT_TEMPLATE_SHA256=" << template_sha << "\n";
    }
    std::cout << "RECORD_COUNT=" << result.records.size() << "\n";
    std::cout << "TRAIN_RECORD_COUNT=" << result.train_records << "\n";
    std::cout << "VALIDATION_RECORD_COUNT=" << result.validation_records << "\n";
    std::cout << "ASSISTANT_MESSAGE_COUNT=" << result.assistant_messages << "\n";
    std::cout << "SUPERVISED_TOKEN_COUNT=" << result.supervised_tokens << "\n";
    std::cout << "TRUNCATED_RECORD_COUNT=" << result.truncated_records << "\n";
    std::cout << "PROMPT_MASK_VIOLATIONS=" << result.prompt_mask_violations << "\n";
    std::cout << "PADDING_MASK_VIOLATIONS=" << result.padding_mask_violations << "\n";
    std::cout << "BOUNDARY_VIOLATIONS=" << result.boundary_violations << "\n";
    std::cout << "CAUSAL_LABEL_VIOLATIONS=" << result.label_violations << "\n";
    std::cout << "TOKEN_RANGE_VIOLATIONS=" << result.token_range_violations << "\n";
    std::cout << "TOKEN_ROUNDTRIP_VIOLATIONS=" << result.roundtrip_violations << "\n";
    std::cout << "ASSISTANT_TEXT_VIOLATIONS=" << result.assistant_text_violations << "\n";
    std::cout << "TRAIN_PACKED_BLOCK_COUNT=" << result.train_blocks.size() << "\n";
    std::cout << "VALIDATION_PACKED_BLOCK_COUNT=" << result.validation_blocks.size() << "\n";
    std::cout << "SHARD_RELOAD_EXACT=" << (result.shard_reload_exact ? 1 : 0) << "\n";
    if (!result.fingerprint.empty()) {
        std::cout << "DATASET_FINGERPRINT=" << result.fingerprint << "\n";
    }
    std::cout << "GPU_MEMORY_BEFORE_MIB=" << gpu_before << "\n";
    std::cout << "GPU_MEMORY_AFTER_MIB=" << gpu_after << "\n";
    std::cout << "GPU_MEMORY_DELTA_MIB="
              << ((gpu_before >= 0 && gpu_after >= 0) ? gpu_after - gpu_before : -1)
              << "\n";

    return result;
}

static int verify_shard_mode(
        const fs::path & shard_path,
        uint64_t seed,
        uint32_t max_sequence_length) {
    bool checksum_mismatch = false;
    try {
        const auto blocks = read_shard(
            shard_path,
            max_sequence_length,
            seed,
            &checksum_mismatch);
        std::cout << "VERIFIED_BLOCK_COUNT=" << blocks.size() << "\n";
        std::cout << "CHECKSUM_MISMATCH=0\n";
        std::cout << "SUBTEST_STATUS=PASS\n";
        std::cout << "FINAL_STATUS=PASS\n";
        return 0;
    } catch (const std::exception & error) {
        std::cout << "VERIFY_ERROR=" << error.what() << "\n";
        std::cout << "CHECKSUM_MISMATCH=" << (checksum_mismatch ? 1 : 0) << "\n";
        std::cout << "SUBTEST_STATUS=FAIL\n";
        std::cout << "FINAL_STATUS=FAIL\n";
        return 1;
    }
}

int main(int argc, char ** argv) {
    try {
        if (argc < 9) {
            std::cerr
                << "usage: test-q1-lora-dataset MODEL CANONICAL OUTPUT_DIR MODE "
                << "SEED MAX_SEQUENCE_LENGTH VALIDATION_PERMILLE SHARD_BLOCK_CAPACITY "
                << "[SHARD_PATH]\n";
            return 2;
        }

        const fs::path model_path = argv[1];
        const fs::path canonical_path = argv[2];
        const fs::path output_dir = argv[3];
        const std::string mode = argv[4];
        const uint64_t seed = std::stoull(argv[5]);
        const uint32_t max_sequence_length =
            static_cast<uint32_t>(std::stoul(argv[6]));
        const uint32_t validation_permille =
            static_cast<uint32_t>(std::stoul(argv[7]));
        const uint32_t shard_block_capacity =
            static_cast<uint32_t>(std::stoul(argv[8]));

        if (
            max_sequence_length < 8 ||
            validation_permille == 0 ||
            validation_permille >= 1000 ||
            shard_block_capacity == 0) {
            throw std::runtime_error("invalid numeric configuration");
        }

        if (mode == "verify_shard") {
            if (argc < 10) {
                throw std::runtime_error("verify_shard requires SHARD_PATH");
            }
            return verify_shard_mode(
                argv[9],
                seed,
                max_sequence_length);
        }

        const bool write_outputs =
            mode == "shard" ||
            mode == "reload" ||
            mode == "batch" ||
            mode == "full";

        pipeline_result result = build_pipeline(
            model_path,
            canonical_path,
            output_dir,
            seed,
            max_sequence_length,
            validation_permille,
            shard_block_capacity,
            write_outputs);

        bool pass = true;

        if (mode == "api") {
            pass =
                !result.records.empty();
        } else if (mode == "tokenize") {
            pass =
                result.token_range_violations == 0 &&
                result.roundtrip_violations == 0;
        } else if (mode == "mask") {
            pass =
                result.prompt_mask_violations == 0 &&
                result.supervised_tokens > 0;
        } else if (mode == "truncation") {
            pass =
                result.truncated_records > 0 &&
                result.supervised_tokens > 0;
        } else if (mode == "split") {
            pass =
                result.train_records > 0 &&
                result.validation_records > 0 &&
                result.train_records + result.validation_records == result.records.size();
        } else if (mode == "packing") {
            pass =
                result.padding_mask_violations == 0 &&
                result.boundary_violations == 0 &&
                !result.train_blocks.empty() &&
                !result.validation_blocks.empty();
        } else if (mode == "shard") {
            pass =
                result.shards.size() >= 2 &&
                std::all_of(
                    result.shards.begin(),
                    result.shards.end(),
                    [&](const shard_info & shard) {
                        return shard.block_count > 0 &&
                               shard.block_count <= shard_block_capacity;
                    }) &&
                !result.fingerprint.empty();
        } else if (mode == "reload") {
            pass = result.shard_reload_exact;
        } else if (mode == "batch") {
            pass =
                result.label_violations == 0 &&
                result.padding_mask_violations == 0 &&
                result.boundary_violations == 0;
        } else if (mode == "full") {
            pass =
                result.token_range_violations == 0 &&
                result.roundtrip_violations == 0 &&
                result.prompt_mask_violations == 0 &&
                result.padding_mask_violations == 0 &&
                result.boundary_violations == 0 &&
                result.label_violations == 0 &&
                result.train_records > 0 &&
                result.validation_records > 0 &&
                result.shard_reload_exact &&
                !result.fingerprint.empty();
        } else {
            throw std::runtime_error("unsupported mode: " + mode);
        }

        std::cout << "SUBTEST_STATUS=" << (pass ? "PASS" : "FAIL") << "\n";
        std::cout << "FINAL_STATUS=" << (pass ? "PASS" : "FAIL") << "\n";
        return pass ? 0 : 1;

    } catch (const std::exception & error) {
        std::cerr << "STEP9_ERROR=" << error.what() << "\n";
        std::cout << "SUBTEST_STATUS=FAIL\n";
        std::cout << "FINAL_STATUS=FAIL\n";
        return 1;
    }
}