"""
On-device bench and completion test runner for llama.cpp (CPU, GPU, NPU backends).

On Android: calls scripts/snapdragon/run.py on the QDC runner host
(script wraps commands in adb shell internally).

On Linux: runs llama-bench directly via run_linux.sh (BASH framework).

Placeholders replaced at artifact creation time by run_qdc_jobs.py:
  <<MODEL_URL>>  Direct URL to the GGUF model file (downloaded on-device)
"""

import os
import subprocess
import sys

import pytest

from utils import (
    BIN_PATH,
    MODEL_DEVICE_PATH,
    PROMPT_DIR,
    push_bundle_if_needed,
    run_adb_command,
    run_snapdragon,
    write_qdc_log,
)

MODEL_URL = "<<MODEL_URL>>"


@pytest.fixture(scope="session", autouse=True)
def install(driver):
    push_bundle_if_needed(f"{BIN_PATH}/llama-cli")
    run_adb_command(f"mkdir -p /data/local/tmp/gguf {PROMPT_DIR}")
    run_adb_command(f"echo 'What is the capital of France?' > {PROMPT_DIR}/bench_prompt.txt")
    check = subprocess.run(
        ["adb", "shell", f"ls {MODEL_DEVICE_PATH}"],
        text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
    )
    if check.returncode != 0:
        run_adb_command(f'curl -L -J --output {MODEL_DEVICE_PATH} "{MODEL_URL}"')


@pytest.mark.parametrize(
    "device",
    [
        pytest.param("none", id="cpu"),
        pytest.param("GPUOpenCL", id="gpu"),
        pytest.param("HTP0", id="npu"),
    ],
)
def test_llama_completion(device):
    args = [
        "llama-completion",
        "-m", MODEL_DEVICE_PATH,
        "-f", f"{PROMPT_DIR}/bench_prompt.txt",
        "-no-cnv",
        "--ctx-size", "8192",
        "-n", "128",
        "--seed", "42",
    ]
    if device == "HTP0":
        args += ["--ubatch-size", "1024"]
    result = run_snapdragon(args, device=device)
    write_qdc_log(f"llama_completion_{device}.log", result.stdout or "")
    assert result.returncode == 0, (
        f"llama-completion {device} failed (exit {result.returncode})"
    )


_DEVICE_LOG_NAME = {"none": "cpu", "GPUOpenCL": "gpu", "HTP0": "htp"}


@pytest.mark.parametrize(
    "device",
    [
        pytest.param("none", id="cpu"),
        pytest.param("GPUOpenCL", id="gpu"),
        pytest.param("HTP0", id="npu"),
    ],
)
def test_llama_bench(device):
    args = [
        "llama-bench",
        "-m", MODEL_DEVICE_PATH,
        "-ngl", "99",
        "-p", "128",
        "-n", "32",
    ]
    if device == "HTP0":
        args += ["--ubatch-size", "1024"]
    result = run_snapdragon(args, device=device)
    write_qdc_log(f"llama_bench_{_DEVICE_LOG_NAME[device]}.log", result.stdout or "")
    assert result.returncode == 0, (
        f"llama-bench {device} failed (exit {result.returncode})"
    )


if __name__ == "__main__":
    ret = pytest.main(["-s", "--junitxml=results.xml", os.path.realpath(__file__)])
    if os.path.exists("results.xml"):
        with open("results.xml") as f:
            write_qdc_log("results.xml", f.read())
    sys.exit(ret)