// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// forge CLI:
//   forge train    --config configs/gpt-25m.json --data data/tinystories --out runs/exp1
//                  [--resume runs/exp1/ckpt_latest.bin] [--backend metal|cpu]
//   forge generate --checkpoint runs/exp1/ckpt_latest.bin --tokenizer data/tinystories/tok4096.model
//                  --prompt "Once upon a time" [--temp 0.8] [--top-k 40] [--max-tokens 256]
//   forge eval     --checkpoint runs/exp1/ckpt_latest.bin --data data/tinystories/val.bin
//                  [--batches 50]
//   forge info     [--config <json>]
#include <Foundation/Foundation.hpp>
#include <Metal/Metal.hpp>

#include "core/device.h"
#include "core/fmodel.h"
#include "nn/config.h"
#include "nn/transformer.h"
#include "ops/metal/metal_ops.h"
#include "ops/ops.h"
#include "tokenizer/bpe.h"
#include "train/checkpoint.h"
#include "train/dataloader.h"
#include "train/trainer.h"

#include <cmath>
#include <cstdio>
#include <filesystem>
#include <fstream>
#include <map>
#include <memory>
#include <random>
#include <sstream>
#include <string>
#include <unordered_set>
#include <vector>

namespace {

void print_usage() {
    std::printf(
        "usage: forge <command> [options]\n"
        "\n"
        "commands:\n"
        "  train     --config <json> --data <dir> --out <dir> [--resume <ckpt>] [--backend metal|cpu]\n"
        "  generate  --checkpoint <ckpt> --tokenizer <model> --prompt <text>\n"
        "            [--temp t] [--top-k k] [--max-tokens n] [--seed s]\n"
        "  eval      --checkpoint <ckpt> --data <val.bin> [--batches n]\n"
        "  export    --checkpoint <ckpt.bin> --out <model.forge> [--dtype f32|f16|bf16]\n"
        "            [--shard-mb n] [--tag label]   (git-style delta repo, zero-copy load)\n"
        "  info      [--config <json>]\n"
        "\n"
        "generate/eval also accept a .forge repo (or one of its manifests) as --checkpoint.\n");
}

std::map<std::string, std::string> parse_flags(int argc, char** argv, int start) {
    std::map<std::string, std::string> flags;
    for (int i = start; i < argc; ++i) {
        std::string arg = argv[i];
        if (arg.rfind("--", 0) != 0) continue;
        std::string key = arg.substr(2);
        if (i + 1 < argc && std::string(argv[i + 1]).rfind("--", 0) != 0) {
            flags[key] = argv[++i];
        } else {
            flags[key] = "true";
        }
    }
    return flags;
}

std::string flag(const std::map<std::string, std::string>& flags, const std::string& key,
                 const std::string& fallback = "") {
    auto it = flags.find(key);
    return it == flags.end() ? fallback : it->second;
}

void select_backend(const std::string& name) {
    if (name == "cpu") {
        forge::ops::set_backend(forge::ops::Backend::CPU);
    } else {
        forge::ops::set_backend(forge::ops::Backend::Metal);
        forge::Device::get(); // fail fast if Metal init is broken
    }
}

int cmd_info(const std::map<std::string, std::string>& flags) {
    forge::Device& dev = forge::Device::get();
    std::printf("device: %s\n", dev.name().c_str());
    std::printf("recommended working set: %.1f GB\n",
                double(dev.recommended_working_set()) / (1024.0 * 1024.0 * 1024.0));
    if (auto path = flag(flags, "config"); !path.empty()) {
        forge::Config cfg = forge::load_config(path);
        const auto& m = cfg.model;
        std::printf("\nmodel: %s\n", m.name.c_str());
        std::printf("  layers=%lld d_model=%lld heads=%lld kv_heads=%lld head_dim=%lld\n",
                    m.n_layers, m.d_model, m.n_heads, m.n_kv_heads, m.head_dim());
        std::printf("  d_ff=%lld vocab=%lld context=%lld\n", m.d_ff, m.vocab_size,
                    m.context_length);
        std::printf("  norm=%s act=%s rope=%d tied=%d\n", m.norm.c_str(),
                    m.activation.c_str(), m.use_rope, m.tied_embeddings);
        std::printf("  parameters: %.2fM\n", double(m.num_params()) / 1e6);
    }
    return 0;
}

int cmd_train(const std::map<std::string, std::string>& flags) {
    const std::string config_path = flag(flags, "config");
    const std::string data_dir = flag(flags, "data");
    const std::string out_dir = flag(flags, "out");
    if (config_path.empty() || data_dir.empty() || out_dir.empty()) {
        print_usage();
        return 1;
    }
    select_backend(flag(flags, "backend", "metal"));

    std::ifstream in(config_path);
    std::stringstream ss;
    ss << in.rdbuf();
    const std::string config_json = ss.str();

    forge::Config cfg = forge::load_config(config_path);
    forge::train::Trainer trainer(cfg, data_dir, out_dir, config_json);
    trainer.train(flag(flags, "resume"));
    return 0;
}

// True when the path is a .forge repo directory or a manifest json inside one.
bool is_forge_repo(const std::string& path) {
    namespace fs = std::filesystem;
    if (fs::is_directory(path)) return fs::exists(fs::path(path) / "manifest-latest.json");
    return path.size() > 5 && path.rfind(".json") == path.size() - 5;
}

// Build the model described by a .forge snapshot and alias its weights.
// f32 repos are ZERO-COPY: parameters point straight into the mmapped,
// page-aligned shards (unified memory — the file cache feeds the GPU).
std::unique_ptr<forge::nn::Transformer> load_forge_model(const std::string& path,
                                                         forge::Config* out_cfg) {
    forge::fmodel::Snapshot snap = forge::fmodel::Snapshot::open(path);
    nlohmann::json j = nlohmann::json::parse(snap.config_json());
    forge::Config cfg;
    if (j.contains("model")) j.at("model").get_to(cfg.model);
    if (j.contains("train")) j.at("train").get_to(cfg.train);
    auto model = std::make_unique<forge::nn::Transformer>(cfg.model, 0);

    std::unordered_set<const void*> seen;
    for (const auto& [name, p] : model->named_parameters()) {
        if (!seen.insert(p.id()).second) continue; // tied param: already aliased
        if (!snap.has(name)) {
            std::fprintf(stderr, "forge: %s missing from %s\n", name.c_str(),
                         path.c_str());
            std::exit(1);
        }
        if (snap.shape(name) != p.value().shape()) {
            std::fprintf(stderr, "forge: shape mismatch for %s\n", name.c_str());
            std::exit(1);
        }
        p.value() = snap.tensor_f32(name);
    }
    if (out_cfg) *out_cfg = cfg;
    std::printf("loaded %s (step %lld)\n", path.c_str(),
                static_cast<long long>(snap.step()));
    return model;
}

int cmd_export(const std::map<std::string, std::string>& flags) {
    const std::string ckpt = flag(flags, "checkpoint");
    const std::string out = flag(flags, "out");
    if (ckpt.empty() || out.empty()) {
        print_usage();
        return 1;
    }
    select_backend(flag(flags, "backend", "metal"));

    const std::string config_json = forge::train::read_checkpoint_config(ckpt);
    nlohmann::json j = nlohmann::json::parse(config_json);
    forge::ModelConfig mc;
    j.at("model").get_to(mc);
    forge::nn::Transformer model(mc, 0);
    forge::train::CheckpointData meta =
        forge::train::load_checkpoint(ckpt, model.named_parameters(), nullptr);

    forge::fmodel::SaveOptions opts;
    const std::string dt = flag(flags, "dtype", "f32");
    opts.dtype = dt == "f16" ? forge::DType::F16
               : dt == "bf16" ? forge::DType::BF16
                              : forge::DType::F32;
    opts.shard_mb = std::stoll(flag(flags, "shard-mb", "95"));
    opts.tag = flag(flags, "tag");
    opts.step = meta.step;
    forge::fmodel::save(out, config_json, model.named_parameters(), opts);
    return 0;
}

int cmd_generate(const std::map<std::string, std::string>& flags) {
    const std::string ckpt = flag(flags, "checkpoint");
    const std::string tok_path = flag(flags, "tokenizer");
    if (ckpt.empty() || tok_path.empty()) {
        print_usage();
        return 1;
    }
    select_backend(flag(flags, "backend", "metal"));
    const float temp = std::stof(flag(flags, "temp", "0.8"));
    const int64_t top_k = std::stoll(flag(flags, "top-k", "40"));
    const int64_t max_tokens = std::stoll(flag(flags, "max-tokens", "256"));
    const uint64_t seed = std::stoull(flag(flags, "seed", "1234"));

    // model config comes from the checkpoint / .forge repo itself
    std::unique_ptr<forge::nn::Transformer> model_ptr;
    forge::ModelConfig mc;
    if (is_forge_repo(ckpt)) {
        forge::Config cfg;
        model_ptr = load_forge_model(ckpt, &cfg);
        mc = cfg.model;
    } else {
        nlohmann::json j =
            nlohmann::json::parse(forge::train::read_checkpoint_config(ckpt));
        j.at("model").get_to(mc);
        model_ptr = std::make_unique<forge::nn::Transformer>(mc, 0);
        forge::train::load_checkpoint(ckpt, model_ptr->named_parameters(), nullptr);
    }
    forge::nn::Transformer& model = *model_ptr;

    forge::tok::BPETokenizer tokenizer;
    tokenizer.load(tok_path);

    std::vector<int32_t> ctx = tokenizer.encode(flag(flags, "prompt", "Once upon a time"));
    if (ctx.empty()) ctx.push_back(0);
    std::printf("%s", tokenizer.decode(ctx).c_str());
    std::fflush(stdout);

    std::mt19937_64 rng(seed);
    forge::NoGrad ng;
    for (int64_t n = 0; n < max_tokens; ++n) {
        // full-context recompute each token (KV cache lands with M5)
        const int64_t T =
            std::min<int64_t>(int64_t(ctx.size()), mc.context_length);
        forge::Tensor ids = forge::Tensor::empty({1, T}, forge::DType::I32);
        for (int64_t t = 0; t < T; ++t)
            ids.data<int32_t>()[t] = ctx[ctx.size() - size_t(T) + size_t(t)];

        forge::Var logits = model.forward(ids); // [T, V]
        if (forge::ops::backend() == forge::ops::Backend::Metal) forge::metal::sync();

        const float* row = logits.value().data<float>() + (T - 1) * mc.vocab_size;
        std::vector<std::pair<float, int32_t>> cand(size_t(mc.vocab_size));
        for (int64_t v = 0; v < mc.vocab_size; ++v) cand[size_t(v)] = {row[v], int32_t(v)};
        const size_t k = size_t(std::min<int64_t>(top_k > 0 ? top_k : mc.vocab_size,
                                                  mc.vocab_size));
        std::partial_sort(cand.begin(), cand.begin() + long(k), cand.end(),
                          [](auto& a, auto& b) { return a.first > b.first; });

        // temperature softmax over the top-k
        float m = cand[0].first;
        double sum = 0.0;
        std::vector<double> probs(k);
        const float tinv = temp > 0.0f ? 1.0f / temp : 1.0f;
        for (size_t i = 0; i < k; ++i) {
            probs[i] = std::exp(double((cand[i].first - m) * tinv));
            sum += probs[i];
        }
        double r = std::uniform_real_distribution<double>(0.0, sum)(rng);
        int32_t next = cand[0].second;
        for (size_t i = 0; i < k; ++i) {
            r -= probs[i];
            if (r <= 0.0) {
                next = cand[i].second;
                break;
            }
        }
        ctx.push_back(next);
        std::printf("%s", tokenizer.decode({next}).c_str());
        std::fflush(stdout);
    }
    std::printf("\n");
    return 0;
}

int cmd_eval(const std::map<std::string, std::string>& flags) {
    const std::string ckpt = flag(flags, "checkpoint");
    const std::string data = flag(flags, "data");
    if (ckpt.empty() || data.empty()) {
        print_usage();
        return 1;
    }
    select_backend(flag(flags, "backend", "metal"));
    const int64_t batches = std::stoll(flag(flags, "batches", "50"));

    std::unique_ptr<forge::nn::Transformer> model_ptr;
    forge::Config cfg;
    if (is_forge_repo(ckpt)) {
        model_ptr = load_forge_model(ckpt, &cfg);
    } else {
        nlohmann::json j =
            nlohmann::json::parse(forge::train::read_checkpoint_config(ckpt));
        if (j.contains("model")) j.at("model").get_to(cfg.model);
        if (j.contains("train")) j.at("train").get_to(cfg.train);
        model_ptr = std::make_unique<forge::nn::Transformer>(cfg.model, 0);
        forge::train::load_checkpoint(ckpt, model_ptr->named_parameters(), nullptr);
    }
    forge::nn::Transformer& model = *model_ptr;

    forge::train::DataLoader loader(data, cfg.model.context_length, 0);
    const int64_t B = cfg.train.batch_size, T = cfg.model.context_length;

    forge::NoGrad ng;
    double total = 0.0;
    for (int64_t i = 0; i < batches; ++i) {
        forge::Tensor ids = forge::Tensor::empty({B, T}, forge::DType::I32);
        forge::Tensor targets = forge::Tensor::empty({B * T}, forge::DType::I32);
        loader.seq_batch(i, ids, targets);
        forge::Var loss = model.loss(ids, targets);
        if (forge::ops::backend() == forge::ops::Backend::Metal) forge::metal::sync();
        total += double(loss.value().data<float>()[0]);
    }
    const double mean = total / double(batches);
    std::printf("val loss: %.4f | ppl: %.2f (%lld batches of %lldx%lld)\n", mean,
                std::exp(mean), static_cast<long long>(batches),
                static_cast<long long>(B), static_cast<long long>(T));
    return 0;
}

} // namespace

int main(int argc, char** argv) {
    NS::AutoreleasePool* pool = NS::AutoreleasePool::alloc()->init();

    int rc = 0;
    if (argc < 2) {
        print_usage();
        rc = 1;
    } else {
        const std::string command = argv[1];
        const auto flags = parse_flags(argc, argv, 2);
        if (command == "info") rc = cmd_info(flags);
        else if (command == "train") rc = cmd_train(flags);
        else if (command == "generate") rc = cmd_generate(flags);
        else if (command == "eval") rc = cmd_eval(flags);
        else if (command == "export") rc = cmd_export(flags);
        else {
            print_usage();
            rc = 1;
        }
    }

    pool->drain();
    return rc;
}