// Yasha Engine CLI — interactive inference with all quantization, diffusion, and config options
// Build: g++ -std=c++17 -mavx2 -mfma -O3 -pthread yasha.cpp main.cpp -o yasha

#include "yasha.h"
#include <iostream>
#include <cstring>
#include <chrono>
#include <iomanip>
#include <unistd.h>

void print_usage() {
    std::cerr << "Yasha Engine — CPU inference for Yasha-8B-Preview\n";
    std::cerr << "Usage: ./yasha --model MODEL_DIR [options]\n";
    std::cerr << "Options:\n";
    std::cerr << "  --model DIR         Path to model directory (required)\n";
    std::cerr << "  --prompt TEXT       Input prompt\n";
    std::cerr << "  --tokens N          Max tokens to generate (default 512)\n";
    std::cerr << "  --temp F            Temperature (default 0.7)\n";
    std::cerr << "  --top-p F           Top-p sampling (default 0.9)\n";
    std::cerr << "  --quant MODE        Quantization: fp32, nf4, q2, q3, auto (default auto)\n";
    std::cerr << "  --diffusion LEVEL   Self-diffusion: 0=off, 1=token, 2=full, 3=regenerate (default 1)\n";
    std::cerr << "  --moe MODE          MoE: full, merged, adaptive (default adaptive)\n";
    std::cerr << "  --fused-qkv         Use fused QKV projection (default on)\n";
    std::cerr << "  --packed-weights    Use SIMD-packed weights (default on)\n";
    std::cerr << "  --threads N         Thread count (default 4)\n";
    std::cerr << "  --interactive       Interactive mode (multi-turn)\n";
    std::cerr << "  --benchmark         Run speed benchmark\n";
    std::cerr << "  --cache-dir DIR     Prompt cache directory\n";
    std::cerr << "  --confidence        Show confidence scores\n";
    std::cerr << "  --stats             Show performance stats\n";
    std::cerr << "  --version           Print version\n";
}

void print_stats(int64_t ns, int tokens) {
    double ms = ns / 1e6;
    double tps = tokens / (ns / 1e9);
    std::cerr << "\n[Stats] " << tokens << " tokens in " << (int)ms << "ms = "
              << std::fixed << std::setprecision(1) << tps << " tok/s\n";
}

int main(int argc, char** argv) {
    if (argc < 2) { print_usage(); return 1; }
    if (argc == 2 && (!strcmp(argv[1], "--help") || !strcmp(argv[1], "-h"))) { print_usage(); return 0; }
    if (argc == 2 && !strcmp(argv[1], "--version")) {
        std::cerr << "Yasha Engine v1.0 — GLA+MoE+Diffusion CPU inference\n";
        return 0;
    }

    std::string model_dir, prompt, cache_dir;
    int n_pred = 512;
    float temp = 0.7f, top_p = 0.9f;
    std::string quant = "auto";
    int diffusion = 1;
    std::string moe_mode = "adaptive";
    bool fused_qkv = true, packed = true, interactive = false, benchmark = false;
    bool show_confidence = false, show_stats = false;
    int n_threads = 4;

    for (int i = 1; i < argc; i++) {
        auto next = [&]() { return (i + 1 < argc) ? std::string(argv[++i]) : std::string(); };
        if (!strcmp(argv[i], "--model")) model_dir = next();
        else if (!strcmp(argv[i], "--prompt")) prompt = next();
        else if (!strcmp(argv[i], "--tokens")) n_pred = std::stoi(next());
        else if (!strcmp(argv[i], "--temp")) temp = std::stof(next());
        else if (!strcmp(argv[i], "--top-p")) top_p = std::stof(next());
        else if (!strcmp(argv[i], "--quant")) quant = next();
        else if (!strcmp(argv[i], "--diffusion")) diffusion = std::stoi(next());
        else if (!strcmp(argv[i], "--moe")) moe_mode = next();
        else if (!strcmp(argv[i], "--fused-qkv")) fused_qkv = true;
        else if (!strcmp(argv[i], "--no-fused-qkv")) fused_qkv = false;
        else if (!strcmp(argv[i], "--packed-weights")) packed = true;
        else if (!strcmp(argv[i], "--no-packed-weights")) packed = false;
        else if (!strcmp(argv[i], "--threads")) n_threads = std::stoi(next());
        else if (!strcmp(argv[i], "--interactive")) interactive = true;
        else if (!strcmp(argv[i], "--benchmark")) benchmark = true;
        else if (!strcmp(argv[i], "--cache-dir")) cache_dir = next();
        else if (!strcmp(argv[i], "--confidence")) show_confidence = true;
        else if (!strcmp(argv[i], "--stats")) show_stats = true;
    }

    if (model_dir.empty()) { std::cerr << "Error: --model required\n"; return 1; }

    // === Auto-detect quantization ===
    YashaConfig cfg;
    if (quant == "auto") {
        long pages = sysconf(_SC_PHYS_PAGES);
        long page_size = sysconf(_SC_PAGE_SIZE);
        long avail_mb = (pages * page_size) / (1024 * 1024);
        std::cerr << "Detected " << avail_mb << "MB available RAM\n";
        if (avail_mb < 2500) { cfg.weight_bits = 2; std::cerr << "Auto-select: Q2 (1.6GB)\n"; }
        else if (avail_mb < 5000) { cfg.weight_bits = 3; std::cerr << "Auto-select: Q3 (2.4GB)\n"; }
        else if (avail_mb < 12000) { cfg.weight_bits = 1; std::cerr << "Auto-select: NF4 (6.5GB)\n"; }
        else { cfg.weight_bits = 0; std::cerr << "Auto-select: FP32 (26GB)\n"; }
    } else if (quant == "fp32") cfg.weight_bits = 0;
    else if (quant == "nf4") cfg.weight_bits = 1;
    else if (quant == "q2") cfg.weight_bits = 2;
    else if (quant == "q3") cfg.weight_bits = 3;
    else { std::cerr << "Unknown quant: " << quant << "\n"; return 1; }

    cfg.self_diffusion_level = (SelfDiffusionLevel)diffusion;
    cfg.merged_experts = (moe_mode == "merged");
    cfg.adaptive_expert = (moe_mode == "adaptive");
    cfg.fused_qkv = fused_qkv;
    cfg.packed_weights = packed;
    cfg.n_threads = n_threads;

    YashaModel model(cfg);
    auto t0 = std::chrono::high_resolution_clock::now();
    if (!model.load(model_dir)) {
        std::cerr << "Failed to load model from " << model_dir << "\n";
        return 1;
    }
    auto t1 = std::chrono::high_resolution_clock::now();
    std::cerr << "Loaded in " << std::chrono::duration_cast<std::chrono::milliseconds>(t1 - t0).count() << "ms\n";

    // Interactive mode
    if (interactive) {
        std::string line;
        std::cerr << "Interactive mode. Type /quit to exit.\n";
        while (true) {
            std::cerr << ">>> ";
            if (!std::getline(std::cin, line) || line == "/quit") break;
            if (line.empty()) continue;
            std::vector<int> tokens;
            for (char c : line) tokens.push_back((int)(unsigned char)c);
            auto t_start = std::chrono::high_resolution_clock::now();
            Tensor out = model.forward_cached(tokens, n_pred, temp, top_p);
            auto t_end = std::chrono::high_resolution_clock::now();
            int gen = out.numel();
            for (int i = 0; i < gen; i++) std::cout << (char)out.data()[i];
            std::cout << std::flush;
            if (show_confidence && gen > 0) {
                float conf = model.score_confidence_ensemble(out.data(), gen);
                std::cerr << "\n[Confidence: " << std::fixed << std::setprecision(2) << conf << "]";
            }
            if (show_stats)
                print_stats(std::chrono::duration_cast<std::chrono::nanoseconds>(t_end - t_start).count(), gen);
        }
        return 0;
    }

    if (benchmark) {
        std::string bench_prompt = prompt.empty() ? "Hello, how are you?" : prompt;
        std::vector<int> tokens;
        for (char c : bench_prompt) tokens.push_back((int)(unsigned char)c);
        std::cerr << "Benchmarking " << n_pred << " tokens...\n";
        auto t_start = std::chrono::high_resolution_clock::now();
        Tensor out = model.forward(tokens, n_pred, temp, top_p);
        auto t_end = std::chrono::high_resolution_clock::now();
        print_stats(std::chrono::duration_cast<std::chrono::nanoseconds>(t_end - t_start).count(), n_pred);
        for (int i = 0; i < out.numel(); i++) std::cout << (char)out.data()[i];
        std::cout << "\n";
        return 0;
    }

    if (prompt.empty()) {
        std::cerr << "Error: --prompt required (or use --interactive)\n";
        return 1;
    }

    std::vector<int> tokens;
    for (char c : prompt) tokens.push_back((int)(unsigned char)c);
    auto t_start = std::chrono::high_resolution_clock::now();
    Tensor out = model.forward(tokens, n_pred, temp, top_p);
    auto t_end = std::chrono::high_resolution_clock::now();
    for (int i = 0; i < out.numel(); i++) std::cout << (char)out.data()[i];
    std::cout << "\n";
    if (show_stats)
        print_stats(std::chrono::duration_cast<std::chrono::nanoseconds>(t_end - t_start).count(), out.numel());
    if (show_confidence && out.numel() > 0) {
        float conf = model.score_confidence_ensemble(out.data(), out.numel());
        std::cerr << "[Confidence: " << std::fixed << std::setprecision(2) << conf << "]\n";
    }
    return 0;
}