import { exec, spawn } from 'child_process';

const delay = ms => new Promise(res => setTimeout(res, ms));

async function runTest(modelArgs, envVars, name) {
    console.log(`\n===========================================`);
    console.log(`Testing: ${name}`);
    console.log(`Env: ${JSON.stringify(envVars)}`);
    console.log(`Args: ${modelArgs}`);
    
    return new Promise(async (resolve) => {
        await new Promise(r => exec('taskkill /F /IM llama-server.exe', r));
        await delay(2000);

        const env = { ...process.env, ...envVars };
        const server = spawn('llama_bin_run\\llama-server.exe', modelArgs.split(' '), {
            detached: true,
            stdio: 'ignore',
            env
        });

        let ready = false;
        
        for (let i = 0; i < 40; i++) {
            try {
                const res = await fetch('http://127.0.0.1:8000/health', { timeout: 2000 });
                if (res.status === 200) {
                    ready = true;
                    break;
                }
            } catch (e) {}
            await delay(3000);
        }

        if (!ready) {
            console.log(`[${name}] FAILED TO BOOT`);
            exec('taskkill /F /IM llama-server.exe');
            resolve({ success: false });
            return;
        }

        console.log(`[${name}] Server Ready! Running speed test...`);
        exec('node scripts/quick_pptest.mjs', (err, stdout, stderr) => {
            console.log(stdout || stderr);
            exec('taskkill /F /IM llama-server.exe');
            resolve({ success: true });
        });
    });
}

async function main() {
    // 1. 122B-A10B: Pure GPU offload (No n-cpu-moe at all)
    // -ngl 999 will offload all 48 layers to the NVIDIA driver (triggering Shared VRAM fallback on Windows)
    const args122B = `--model models\\Q4_K_M\\Qwen3.5-122B-A10B-Q4_K_M-00001-of-00003.gguf -ngl 999 -c 32768 -np 1 -fa on --cache-type-k q4_0 --cache-type-v q4_0 -ub 512 -b 2048 -t 6 -tb 6 --prio 3 --fit off --port 8000 --host 0.0.0.0`;
    await runTest(args122B, {}, "122B-A10B: Pure GPU (NVIDIA Shared Memory Fallback)");

    // 2. 35B-A3B: Pure GPU tuning to hit 70 t/s
    // Base configuration from previous full-gpu run:
    const args35B = `--model models\\Qwen3.5-35B-A3B-Q4_K_M.gguf -ngl 999 -c 262144 -np 1 -fa on --cache-type-k q4_0 --cache-type-v q4_0 -ub 128 -b 512 -t 6 -tb 6 --prio 3 --fit off --port 8000 --host 0.0.0.0`;
    
    // We already got ~64 t/s basically. 
    // Let's try MMQ for custom matrix multiplication which is often faster on Ampere for low batch generation
    await runTest(args35B, { GGML_CUDA_FORCE_MMQ: "1" }, "35B-A3B: Force MMQ = 1");
    
    // Try increasing threads to 12 just in case
    const args35B_t12 = args35B.replace("-t 6 -tb 6", "-t 12 -tb 12");
    await runTest(args35B_t12, { GGML_CUDA_FORCE_MMQ: "1" }, "35B-A3B: Threads 12 + MMQ");

    console.log("\nALL TESTS COMPLETED");
}

main();