#pragma once

#include <cstdlib>
#include <fstream>
#include <span>
#include <stdexcept>
#include <string>
#include <vector>
#include <array>
#include <cstdint>
#include "ling3/router.h"

namespace ling3 {
// Explicit opt-in diagnostic dumps. The caller creates the directory, and
// should use a short teacher-forced sequence; never enable during benchmarks.
inline thread_local std::size_t numeric_position = 0;
// Optional real attention inputs for the standalone NPU feasibility probe.
// No dumps during ordinary inference; enable only for a short diagnostic run.
inline void NumericMlaCapture(const std::string & name, std::size_t position,
    std::size_t rows, std::span<const float> q, std::span<const std::uint16_t> k,
    std::span<const std::uint16_t> v) {
    static const char * directory = std::getenv("LING3_MLA_CAPTURE");
    if (!directory) return;
    static const std::size_t capture_position = std::getenv("LING3_MLA_CAPTURE_POSITION")
        ? std::stoul(std::getenv("LING3_MLA_CAPTURE_POSITION")) : 384;
    if (position != capture_position) return;
    const std::array<std::uint32_t, 4> header {
        0x31414c4d, static_cast<std::uint32_t>(rows),
        static_cast<std::uint32_t>(position + rows), 16};
    const auto path = std::string(directory) + "/" + name + ".mla";
    std::ofstream file(path, std::ios::binary);
    file.write(reinterpret_cast<const char *>(header.data()), sizeof(header));
    file.write(reinterpret_cast<const char *>(q.data()), q.size_bytes());
    file.write(reinterpret_cast<const char *>(k.data()), k.size_bytes());
    file.write(reinterpret_cast<const char *>(v.data()), v.size_bytes());
    if (!file) throw std::runtime_error("failed MLA capture: " + path);
}
inline void NumericDump(const std::string & name, std::span<const float> values) {
    static const char * directory = std::getenv("LING3_NUMERIC_TRACE");
    if (!directory) return;
    const auto path = std::string(directory) + "/" + std::to_string(numeric_position) + "_" + name + ".f32";
    std::ofstream file(path, std::ios::binary);
    file.write(reinterpret_cast<const char *>(values.data()), values.size_bytes());
    if (!file) throw std::runtime_error("failed numerical trace: " + path);
}
inline void NumericRoutes(const std::string & name, std::span<const Route> routes) {
    if (!std::getenv("LING3_NUMERIC_TRACE")) return;
    std::vector<float> values;
    for (const auto & route : routes) {
        for (const auto id : route.experts) values.push_back(static_cast<float>(id));
        for (const auto weight : route.weights) values.push_back(weight);
    }
    NumericDump(name, values);
}
} // namespace ling3