#include "llama.h"
#include "llama-adapter.h"
#include "llama-model.h"
#include "ggml.h"
#include "ggml-backend.h"

#include <array>
#include <cstdint>
#include <cstdlib>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>

namespace fs = std::filesystem;

struct target_spec {
    const char * name;
    int64_t K;
    int64_t M;
    int64_t rank;
};

static const std::array<target_spec, 3> specs = {{
    {"blk.0.ssm_alpha.weight", 5120, 48, 4},
    {"blk.11.attn_k.weight", 5120, 1024, 4},
    {"blk.0.ffn_down.weight", 17408, 5120, 4},
}};

static std::vector<float> tensor_f32(ggml_tensor * tensor) {
    if (!tensor || tensor->type != GGML_TYPE_F32) {
        throw std::runtime_error("native loader tensor is not F32");
    }
    const size_t count = static_cast<size_t>(ggml_nelements(tensor));
    std::vector<float> values(count, 0.0f);
    ggml_backend_tensor_get(tensor, values.data(), 0, count*sizeof(float));
    return values;
}

static void write_values(const fs::path & path, const std::vector<float> & values) {
    std::ofstream output(path, std::ios::binary | std::ios::trunc);
    output.write(reinterpret_cast<const char *>(values.data()), values.size()*sizeof(float));
    if (!output) {
        throw std::runtime_error("failed writing native tensor dump");
    }
}

int main(int argc, char ** argv) {
    if (argc != 4) {
        std::cerr << "usage: test-q1-lora-roundtrip MODEL ADAPTER OUTPUT_DIR\n";
        return 2;
    }

    try {
        unsetenv("PRISM_Q1_LORA_TRAINING");
        unsetenv("PRISM_Q1_LORA_UNFUSED_GDN");
        unsetenv("PRISM_Q1_LORA_TRAINING_GENERIC_SSM_CONV");
        unsetenv("PRISM_Q1_LORA_TRAINING_NO_KV_CACHE");
        unsetenv("PRISM_STEP10_FORCE_OPT_BACKWARD");
        unsetenv("PRISM_STEP10_EXTERNAL_GRAD_ACCUM");

        const fs::path output_dir = argv[3];
        fs::create_directories(output_dir);

        llama_backend_init();
        llama_model_params params = llama_model_default_params();
        params.n_gpu_layers = 999;
        params.use_mmap = true;
        params.check_tensors = true;

        llama_model * model = llama_model_load_from_file(argv[1], params);
        if (!model) {
            throw std::runtime_error("failed loading base model");
        }

        llama_adapter_lora * adapter = llama_adapter_lora_init(model, argv[2]);
        if (!adapter) {
            llama_model_free(model);
            throw std::runtime_error("failed loading adapter");
        }

        size_t pair_index = 0;
        for (size_t target_index = 0; target_index < specs.size(); ++target_index) {
            const target_spec & spec = specs[target_index];
            const auto found = adapter->ab_map.find(spec.name);
            if (found == adapter->ab_map.end()) {
                throw std::runtime_error(std::string("native loader missing target: ") + spec.name);
            }

            for (const auto & item : std::array<std::pair<const char *, ggml_tensor *>, 2>{{
                    {"a", found->second.a}, {"b", found->second.b}}}) {
                ggml_tensor * tensor = item.second;
                const std::vector<float> values = tensor_f32(tensor);
                const fs::path path = output_dir /
                    ("target_" + std::to_string(target_index) + "_" + item.first + ".bin");
                write_values(path, values);

                std::cout << "NATIVE_TENSOR"
                          << " pair_index=" << pair_index
                          << " target_index=" << target_index
                          << " suffix=" << item.first
                          << " name='" << spec.name << ".lora_" << item.first << "'"
                          << " dtype=" << ggml_type_name(tensor->type)
                          << " ne0=" << tensor->ne[0]
                          << " ne1=" << tensor->ne[1]
                          << " count=" << values.size()
                          << " file='" << path.string() << "'\n";
                ++pair_index;
            }
        }

        llama_adapter_lora_free(adapter);
        llama_model_free(model);
        llama_backend_free();

        std::cout << "NATIVE_TENSOR_COUNT=" << pair_index << "\n";
        std::cout << "FINAL_STATUS=PASS\n";
        return 0;

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