import { describe, expect, test } from "bun:test";
import type {
  Accelerator,
  HostProfile,
  LaunchRequest,
  EngineRuntimeKind,
  ServingOptions,
} from "../src/modules/compute/contracts";
import { ENGINE_IDS } from "../src/modules/compute/contracts";
import { deviceEnvironment, dockerFlagsFor } from "../src/modules/compute/engines/devices";
import { engineSpec, planLaunch, supportsRuntime } from "../src/modules/compute/engines/registry";
import { mergeArguments } from "../src/modules/compute/engines/shared";

const host = (overrides: Partial<HostProfile> = {}): HostProfile => ({
  nodeId: "self",
  platform: "linux",
  arch: "x64",
  accelerator: "cuda",
  unifiedMemory: false,
  wsl: false,
  docker: true,
  dockerGpu: true,
  deviceCount: 2,
  ...overrides,
});

const options = (overrides: Partial<ServingOptions> = {}): ServingOptions => ({
  tensorParallel: 1,
  pipelineParallel: 1,
  maxContextLength: 8192,
  memoryFraction: 0.9,
  maxConcurrentRequests: 64,
  kvCacheDtype: "auto",
  dtype: null,
  quantization: null,
  trustRemoteCode: false,
  toolCallParser: null,
  reasoningParser: null,
  ...overrides,
});

const request = (overrides: Partial<LaunchRequest> = {}): LaunchRequest => ({
  engine: "vllm",
  host: host(),
  runtime: "process",
  devices: ["GPU-aaa", "GPU-bbb"],
  port: 8000,
  modelPath: "/models/qwen",
  servedModelName: "qwen",
  options: options(),
  extraArgs: [],
  env: {},
  binary: "/venv/bin/vllm",
  ...overrides,
});

describe("plan() is pure", () => {
  test("same request produces identical plans", () => {
    const input = request();
    expect(planLaunch(input)).toEqual(planLaunch(input));
  });

  test("every engine emits the model, the port and the devices", () => {
    for (const id of ENGINE_IDS) {
      const spec = engineSpec(id);
      const plan = planLaunch(
        request({ engine: id, host: host({ accelerator: "cuda" }), binary: spec.defaultBinary }),
      );
      expect(plan.argv).toContain("/models/qwen");
      expect(plan.argv).toContain("8000");
      expect(plan.devices).toEqual(["GPU-aaa", "GPU-bbb"]);
      expect(plan.health.path.startsWith("/")).toBe(true);
    }
  });
});

describe("tuning knobs", () => {
  test("vLLM and SGLang spell the same knobs differently", () => {
    const tuned = options({
      tensorParallel: 4,
      maxContextLength: 32768,
      memoryFraction: 0.92,
      toolCallParser: "hermes",
    });
    const vllmArgs = planLaunch(request({ engine: "vllm", options: tuned })).argv;
    const sglangArgs = planLaunch(
      request({ engine: "sglang", options: tuned, port: 30000 }),
    ).argv;

    expect(vllmArgs).toContain("--tensor-parallel-size");
    expect(vllmArgs).toContain("--max-model-len");
    expect(vllmArgs).toContain("--gpu-memory-utilization");
    // vLLM needs the companion flag next to the parser; SGLang does not.
    expect(vllmArgs).toContain("--enable-auto-tool-choice");

    expect(sglangArgs).toContain("--context-length");
    expect(sglangArgs).toContain("--mem-fraction-static");
    expect(sglangArgs).not.toContain("--enable-auto-tool-choice");
    expect(sglangArgs).not.toContain("--max-model-len");
  });

  test("parallelism of 1 is omitted, above 1 is emitted", () => {
    expect(planLaunch(request({ options: options({ tensorParallel: 1 }) })).argv).not.toContain(
      "--tensor-parallel-size",
    );
    expect(planLaunch(request({ options: options({ tensorParallel: 2 }) })).argv).toContain(
      "--tensor-parallel-size",
    );
  });

  test('kvCacheDtype "auto" is treated as unset', () => {
    expect(planLaunch(request({ options: options({ kvCacheDtype: "auto" }) })).argv).not.toContain(
      "--kv-cache-dtype",
    );
    expect(planLaunch(request({ options: options({ kvCacheDtype: "fp8" }) })).argv).toContain(
      "--kv-cache-dtype",
    );
  });

  test("knobs an engine has no equivalent for are dropped, not approximated", () => {
    const argv = planLaunch(
      request({
        engine: "llamacpp",
        binary: "llama-server",
        options: options({ tensorParallel: 4, kvCacheDtype: "fp8", quantization: "awq" }),
      }),
    ).argv;
    expect(argv).not.toContain("--tensor-parallel-size");
    expect(argv).not.toContain("--split-mode");
    expect(argv).not.toContain("--cache-type-k");
    expect(argv).not.toContain("--quantization");
    expect(argv).toContain("--ctx-size");
  });
});

describe("extraArgs override the base flags", () => {
  test("a repeated flag appears once, with the recipe's value", () => {
    const argv = planLaunch(
      request({ options: options({ maxContextLength: 8192 }), extraArgs: ["--max-model-len", "4096"] }),
    ).argv;
    expect(argv.filter((token) => token === "--max-model-len")).toHaveLength(1);
    expect(argv[argv.indexOf("--max-model-len") + 1]).toBe("4096");
    expect(argv).not.toContain("8192");
  });

  test("mergeArguments drops the base flag and its value", () => {
    expect(mergeArguments(["--a", "1", "--b", "2"], ["--a", "9"])).toEqual(["--b", "2", "--a", "9"]);
  });

  test("mergeArguments handles boolean flags and --flag=value form", () => {
    expect(mergeArguments(["--verbose", "--b", "2"], ["--verbose"])).toEqual(["--b", "2", "--verbose"]);
    expect(mergeArguments(["--a=1", "--b", "2"], ["--a", "9"])).toEqual(["--b", "2", "--a", "9"]);
  });

  test("SGLang's --enable-metrics default can be overridden", () => {
    const argv = planLaunch(
      request({ engine: "sglang", port: 30000, extraArgs: ["--enable-metrics", "false"] }),
    ).argv;
    expect(argv.filter((token) => token === "--enable-metrics")).toHaveLength(1);
  });
});

describe("docker vs process", () => {
  test("a container gets the image, the mount and 0.0.0.0; a process gets the binary", () => {
    const asProcess = planLaunch(request({ runtime: "process" }));
    const asDocker = planLaunch(request({ runtime: "docker" }));

    expect(asProcess.argv[0]).toBe("/venv/bin/vllm");
    expect(asProcess.image).toBeUndefined();
    expect(asProcess.mounts).toEqual([]);
    expect(asProcess.argv[asProcess.argv.indexOf("--host") + 1]).toBe("127.0.0.1");

    expect(asDocker.argv[0]).not.toBe("/venv/bin/vllm");
    expect(asDocker.image).toBe("vllm/vllm-openai:latest");
    expect(asDocker.mounts).toEqual([{ from: "/models/qwen", to: "/models", readOnly: true }]);
    expect(asDocker.argv[asDocker.argv.indexOf("--host") + 1]).toBe("0.0.0.0");
    // The container sees the model at the mount point, not the host path.
    expect(asDocker.argv).toContain("/models");
  });
});

describe("device translation", () => {
  const devices = ["GPU-aaa", "GPU-bbb"];

  test("each accelerator selects devices its own way", () => {
    expect(deviceEnvironment("cuda", devices)).toEqual({ CUDA_VISIBLE_DEVICES: "GPU-aaa,GPU-bbb" });
    expect(deviceEnvironment("rocm", ["gpu:0", "gpu:1"])).toEqual({
      HIP_VISIBLE_DEVICES: "0,1",
      ROCR_VISIBLE_DEVICES: "0,1",
    });
    expect(deviceEnvironment("xpu", ["xpu:0"])).toEqual({ ONEAPI_DEVICE_SELECTOR: "level_zero:0" });
    // Metal exposes no device selection and CPU has nothing to select.
    expect(deviceEnvironment("metal", ["metal:0"])).toEqual({});
    expect(deviceEnvironment("cpu", ["cpu:0"])).toEqual({});
  });

  test("no devices means no environment, for every accelerator", () => {
    const accelerators: Accelerator[] = ["cuda", "rocm", "metal", "xpu", "cpu"];
    for (const accelerator of accelerators) {
      expect(deviceEnvironment(accelerator, [])).toEqual({});
      expect(dockerFlagsFor(accelerator, []).args).toEqual([]);
    }
  });

  test("docker flags differ from process env", () => {
    expect(dockerFlagsFor("cuda", devices).args).toEqual(["--gpus", "device=GPU-aaa,GPU-bbb"]);
    expect(dockerFlagsFor("rocm", devices).args).toContain("/dev/kfd");
    expect(dockerFlagsFor("rocm", devices).groupAdd).toEqual(["video", "render"]);
  });

  test("planLaunch folds device selection in exactly once", () => {
    const plan = planLaunch(request());
    expect(plan.env["CUDA_VISIBLE_DEVICES"]).toBe("GPU-aaa,GPU-bbb");
    // The engine must not have written it itself.
    expect(engineSpec("vllm").plan(request()).env["CUDA_VISIBLE_DEVICES"]).toBeUndefined();
  });
});

describe("supports() gates by host", () => {
  const cases: readonly {
    readonly engine: (typeof ENGINE_IDS)[number];
    readonly host: HostProfile;
    readonly ok: boolean;
    readonly runtime?: EngineRuntimeKind;
  }[] = [
    { engine: "vllm", host: host({ accelerator: "cuda" }), ok: true, runtime: "process" },
    { engine: "vllm", host: host({ platform: "darwin", accelerator: "metal" }), ok: false },
    { engine: "vllm", host: host({ platform: "win32", accelerator: "cuda", wsl: false }), ok: false },
    { engine: "vllm", host: host({ platform: "win32", accelerator: "cuda", wsl: true }), ok: true },
    { engine: "vllm", host: host({ accelerator: "rocm" }), ok: true, runtime: "docker" },
    { engine: "sglang", host: host({ accelerator: "rocm" }), ok: false },
    { engine: "mlx", host: host({ platform: "darwin", arch: "arm64", accelerator: "metal" }), ok: true, runtime: "process" },
    { engine: "mlx", host: host({ platform: "linux" }), ok: false },
    { engine: "mlx", host: host({ platform: "darwin", arch: "x64", accelerator: "metal" }), ok: false },
    { engine: "llamacpp", host: host({ platform: "darwin", accelerator: "metal" }), ok: true, runtime: "process" },
    { engine: "llamacpp", host: host({ platform: "win32", accelerator: "rocm" }), ok: true, runtime: "process" },
    { engine: "exllamav3", host: host({ accelerator: "cuda" }), ok: true, runtime: "process" },
    { engine: "exllamav3", host: host({ accelerator: "rocm" }), ok: false },
  ];

  for (const testCase of cases) {
    const label = `${testCase.engine} on ${testCase.host.platform}/${testCase.host.accelerator}${testCase.host.wsl ? "+wsl" : ""}/${testCase.host.arch}`;
    test(`${label} -> ${testCase.ok ? "supported" : "refused"}`, () => {
      const support = engineSpec(testCase.engine).supports(testCase.host);
      expect(support.ok).toBe(testCase.ok);
      if (!support.ok) {
        // A refusal must say why, so the UI never has to invent a reason.
        expect(support.reason.length).toBeGreaterThan(10);
        return;
      }
      if (testCase.runtime) {
        expect(supportsRuntime(testCase.engine, testCase.host, testCase.runtime)).toBe(true);
      }
    });
  }

  test("vLLM on ROCm without docker GPU passthrough is refused", () => {
    const support = engineSpec("vllm").supports(host({ accelerator: "rocm", dockerGpu: false }));
    expect(support.ok).toBe(false);
  });

  test("MLX never offers docker — macOS has no Metal passthrough", () => {
    const support = engineSpec("mlx").supports(
      host({ platform: "darwin", arch: "arm64", accelerator: "metal", dockerGpu: true }),
    );
    expect(support.ok && support.runtimes).toEqual(["process"]);
  });
});

describe("DGX Spark", () => {
  const spark = host({ platform: "linux", arch: "arm64", accelerator: "cuda", unifiedMemory: true });

  test("is not a special case — CUDA engines run natively", () => {
    expect(engineSpec("vllm").supports(spark).ok).toBe(true);
    expect(engineSpec("llamacpp").supports(spark).ok).toBe(true);
    expect(planLaunch(request({ host: spark })).env["CUDA_VISIBLE_DEVICES"]).toBe("GPU-aaa,GPU-bbb");
  });
});