# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""Unit tests for the llama-server pass-through args validator.

The validator is the boundary between user CLI/HTTP input and the
llama-server subprocess. These tests pin denylist behaviour so it doesn't
regress when new managed flags are added.
"""

from __future__ import annotations

import importlib.util
import re
from pathlib import Path

import pytest

# Load llama_server_args.py directly to avoid dragging in the full backend
# chain via core/inference/__init__.py. The validator is dependency-free.
_LSA_PATH = Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_server_args.py"
_spec = importlib.util.spec_from_file_location("_lsa_test_only", _LSA_PATH)
_lsa = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(_lsa)
is_managed_flag = _lsa.is_managed_flag
parse_cache_override = _lsa.parse_cache_override
parse_ctx_override = _lsa.parse_ctx_override
resolve_cache_type_kv = _lsa.resolve_cache_type_kv
strip_shadowing_flags = _lsa.strip_shadowing_flags
validate_extra_args = _lsa.validate_extra_args


# ── Pass-through (allowed) ───────────────────────────────────────────


@pytest.mark.parametrize(
    "args",
    [
        # Sampling
        ["--top-k", "20"],
        ["--top-p", "0.9", "--min-p", "0.05"],
        ["--seed", "-1"],  # negative value, not a flag
        ["--temp", "0.0"],
        ["--repeat-penalty", "1.05"],
        ["--mirostat", "2", "--mirostat-lr", "0.1"],
        ["--xtc-probability", "0.05", "--xtc-threshold", "0.1"],
        ["--dry-multiplier", "0.5"],
        # Tier-2 knobs that map to LoadRequest fields
        ["--cache-type-k", "q8_0"],
        ["--cache-type-v", "q8_0"],
        ["--chat-template-file", "/tmp/tpl.jinja"],
        ["--chat-template-kwargs", '{"reasoning_effort":"high"}'],
        ["--spec-type", "ngram-mod"],
        ["--spec-default"],
        # MTP path (llama.cpp #22673).
        ["--spec-type", "draft-mtp"],
        ["--spec-type", "draft-mtp", "--spec-draft-n-max", "6"],
        [
            "--spec-type",
            "ngram-mod,draft-mtp",
            "--spec-draft-n-max",
            "3",
            "--spec-ngram-mod-n-match",
            "24",
            "--spec-ngram-mod-n-min",
            "48",
            "--spec-ngram-mod-n-max",
            "64",
        ],
        # Reasoning controls
        ["--reasoning-format", "deepseek"],
        ["-rea", "auto"],
        # Soft-managed: user flags last-wins over Studio's auto-set version.
        # --parallel / -np / --n-parallel are hard-denied (KV-cache + slot
        # count would desync); use `unsloth studio run --parallel N` instead.
        ["-c", "131072"],
        ["--ctx-size", "8192"],
        ["--flash-attn", "off"],
        ["-fa", "on"],
        ["--no-context-shift"],
        ["--context-shift"],
        ["--jinja"],
        ["--no-jinja"],
        ["-ngl", "-1"],
        ["--gpu-layers", "32"],
        ["-t", "16"],
        ["--threads", "32"],
        ["-fit", "off"],
        ["--fit", "on"],
        ["--fit-ctx", "8192"],
    ],
)
def test_pass_through_allowed(args):
    assert validate_extra_args(args) == args


def test_none_returns_empty_list():
    assert validate_extra_args(None) == []


def test_empty_list_returns_empty_list():
    assert validate_extra_args([]) == []


def test_value_with_equals_form_passes_through():
    assert validate_extra_args(["--top-k=20"]) == ["--top-k=20"]


def test_non_flag_token_passes_through():
    # Bare positionals are passed through; llama-server can reject them.
    assert validate_extra_args(["foo"]) == ["foo"]


# ── Denylist (rejected) ──────────────────────────────────────────────


@pytest.mark.parametrize(
    "denied",
    [
        # Parallel slots -- owned by the typer --parallel flag.
        "-np",
        "--parallel",
        "--n-parallel",
        # Model identity (every alias; bumping llama.cpp must keep every
        # form rejected, not just the long one).
        "-m",
        "--model",
        "-mu",
        "--model-url",
        "-dr",
        "--docker-repo",
        "-hf",
        "-hfr",
        "--hf-repo",
        "-hff",
        "--hf-file",
        "-hfv",
        "-hfrv",
        "--hf-repo-v",
        "-hffv",
        "--hf-file-v",
        "-hft",
        "--hf-token",
        "-mm",
        "--mmproj",
        "-mmu",
        "--mmproj-url",
        # Networking (Studio binds + proxies)
        "--host",
        "--port",
        "--path",
        "--api-prefix",
        "--reuse-port",
        # Auth / TLS
        "--api-key",
        "--api-key-file",
        "--ssl-key-file",
        "--ssl-cert-file",
        # Single-model server (legacy --webui + current --ui group)
        "--webui",
        "--no-webui",
        "--ui",
        "--no-ui",
        "--ui-config",
        "--ui-config-file",
        "--ui-mcp-proxy",
        "--no-ui-mcp-proxy",
        "--models-dir",
        "--models-preset",
        "--models-max",
        "--models-autoload",
        "--no-models-autoload",
        # Server-mode flips: --embedding / --rerank restrict llama-server to
        # those endpoints and break Studio's chat hop.
        "--embedding",
        "--embeddings",
        "--rerank",
        "--reranking",
        # llama-server's own --tools clashes with Studio's tool policy.
        "--tools",
    ],
)
def test_denylist_rejects_all_aliases(denied):
    with pytest.raises(ValueError, match = denied):
        validate_extra_args([denied, "value"])


@pytest.mark.parametrize(
    "args,offending",
    [
        # Pass-through --parallel would last-wins-override the real slot
        # count while Studio's KV-cache fit + llama_parallel_slots stay at
        # the typer value -- plan vs. process disagree.
        (["--parallel", "8"], "--parallel"),
        (["--parallel=8"], "--parallel"),
        (["--n-parallel", "16"], "--n-parallel"),
        (["--n-parallel=16"], "--n-parallel"),
        (["-np", "32"], "-np"),
        # Attached short form: Click clusters it CLI-side; HTTP /load with
        # `["-np8"]` must still resolve to managed.
        (["-np8"], "-np"),
        (["-np64"], "-np"),
        # Out-of-range values that would bypass the typer 1..64 guard.
        (["--parallel", "999"], "--parallel"),
        (["-np", "0"], "-np"),
        (["-np999"], "-np"),
        # Signed attached forms; `-np-1` must not slip past.
        (["-np-1"], "-np"),
        (["-np+1"], "-np"),
    ],
)
def test_parallel_flags_are_managed(args, offending):
    with pytest.raises(ValueError, match = re.escape(offending)):
        validate_extra_args(args)


def test_denylist_rejects_equals_form():
    with pytest.raises(ValueError, match = "--port"):
        validate_extra_args(["--port=9000"])


@pytest.mark.parametrize(
    "padded",
    [" --parallel", "--parallel ", "\t--parallel", "  -np", "-np \n", "-np\t"],
)
def test_denylist_rejects_whitespace_padded_forms(padded):
    # `_flag_name` trims whitespace before lookup; else a trailing space
    # could slip a managed flag past the boundary.
    with pytest.raises(ValueError, match = "parallel|np"):
        validate_extra_args([padded, "8"])


@pytest.mark.parametrize(
    "attached",
    ["-np8x", "-np-1foo", "-np+1bar", "-np9zzz"],
)
def test_denylist_rejects_np_with_digit_prefix_and_junk(attached):
    # Backend `_flag_name` must classify the same forms the CLI rewriter
    # expands, else HTTP /load could smuggle `-np8x` through.
    with pytest.raises(ValueError, match = "np"):
        validate_extra_args([attached])


def test_denylist_rejects_short_form_when_long_is_denied():
    # `-m` is the short form of --model; rejecting only the long form
    # would leave a trivial bypass.
    with pytest.raises(ValueError, match = "-m"):
        validate_extra_args(["-m", "/some/other/path.gguf"])


def test_denylist_message_names_offending_flag():
    with pytest.raises(ValueError) as excinfo:
        validate_extra_args(["--top-k", "20", "--api-key", "secret"])
    assert "--api-key" in str(excinfo.value)


def test_first_denied_flag_short_circuits():
    # Validation stops at the first denied flag; the message names it.
    with pytest.raises(ValueError, match = "--port"):
        validate_extra_args(["--port", "1", "--host", "x"])


# ── Numeric values that look flag-ish ─────────────────────────────────


@pytest.mark.parametrize("value", ["-1", "-0.5", "-42", "-.5"])
def test_negative_number_value_is_not_flag(value):
    # `--seed -1`: the -1 is a value, not a flag.
    assert validate_extra_args(["--seed", value]) == ["--seed", value]


# ── is_managed_flag helper ───────────────────────────────────────────


def test_is_managed_flag_true_for_denied():
    assert is_managed_flag("--port") is True
    assert is_managed_flag("--api-key") is True
    assert is_managed_flag("-m") is True
    assert is_managed_flag("--model") is True
    # Parallel slots owned by the typer --parallel flag.
    assert is_managed_flag("--parallel") is True
    assert is_managed_flag("--n-parallel") is True
    assert is_managed_flag("-np") is True
    # Normalised forms must classify like the canonical token so
    # is_managed_flag filtering stays in sync with validate_extra_args.
    assert is_managed_flag("-np8") is True
    assert is_managed_flag("--parallel=8") is True
    assert is_managed_flag("--port=9000") is True


def test_is_managed_flag_false_for_pass_through():
    assert is_managed_flag("--top-k") is False
    assert is_managed_flag("--cache-type-k") is False
    assert is_managed_flag("--chat-template-file") is False
    # Soft-managed flags pass through (last-wins override)
    assert is_managed_flag("-c") is False
    assert is_managed_flag("--ctx-size") is False
    assert is_managed_flag("--flash-attn") is False
    assert is_managed_flag("-ngl") is False
    assert is_managed_flag("--threads") is False


# ── strip_shadowing_flags ─────────────────────────────────────────────


def test_strip_shadowing_flags_drops_context_when_requested():
    out = strip_shadowing_flags(
        ["-c", "4096", "--top-k", "20"],
        strip_context = True,
        strip_cache = False,
        strip_spec = False,
        strip_template = False,
    )
    assert out == ["--top-k", "20"]


def test_strip_shadowing_flags_keeps_context_when_not_requested():
    out = strip_shadowing_flags(
        ["-c", "4096", "--top-k", "20"],
        strip_context = False,
        strip_cache = False,
        strip_spec = False,
        strip_template = False,
    )
    assert out == ["-c", "4096", "--top-k", "20"]


def test_strip_shadowing_flags_keeps_chat_template_when_template_disabled():
    # No chat_template_override supplied; inherited
    # --chat-template-file must survive.
    out = strip_shadowing_flags(
        ["--chat-template-file", "/tmp/custom.jinja", "--top-k", "20"],
        strip_context = True,
        strip_cache = True,
        strip_spec = True,
        strip_template = False,
    )
    assert out == ["--chat-template-file", "/tmp/custom.jinja", "--top-k", "20"]


def test_strip_shadowing_flags_drops_template_when_requested():
    out = strip_shadowing_flags(
        ["--chat-template-file", "/tmp/custom.jinja", "--top-k", "20"],
        strip_template = True,
    )
    assert out == ["--top-k", "20"]


def test_strip_shadowing_flags_keeps_cache_when_cache_disabled():
    out = strip_shadowing_flags(
        ["--cache-type-k", "q8_0", "--cache-type-v", "q8_0", "--top-k", "20"],
        strip_cache = False,
    )
    assert out == ["--cache-type-k", "q8_0", "--cache-type-v", "q8_0", "--top-k", "20"]


def test_strip_shadowing_flags_keeps_spec_when_spec_disabled():
    out = strip_shadowing_flags(
        ["--spec-type", "ngram-mod", "--draft-min", "48", "--top-k", "20"],
        strip_spec = False,
    )
    assert out == ["--spec-type", "ngram-mod", "--draft-min", "48", "--top-k", "20"]


def test_strip_shadowing_flags_drops_mtp_flags_when_requested():
    # MTP / draft-mtp flags must drop when speculative_type re-applies.
    out = strip_shadowing_flags(
        [
            "--spec-type",
            "draft-mtp",
            "--spec-draft-n-max",
            "6",
            "--spec-ngram-mod-n-match",
            "24",
            "--spec-ngram-mod-n-min",
            "48",
            "--spec-ngram-mod-n-max",
            "6",
            "--top-k",
            "20",
        ],
        strip_spec = True,
    )
    assert out == ["--top-k", "20"]


def test_is_managed_flag_false_for_mtp_pass_through():
    assert is_managed_flag("--spec-draft-n-max") is False
    assert is_managed_flag("--spec-ngram-mod-n-match") is False
    assert is_managed_flag("--spec-ngram-mod-n-min") is False
    assert is_managed_flag("--spec-ngram-mod-n-max") is False


# ── parse_ctx_override ───────────────────────────────────────────────


@pytest.mark.parametrize(
    "args,expected",
    [
        (None, None),
        ([], None),
        (["--top-k", "20"], None),
        (["--ctx-size", "128000"], 128000),
        (["--ctx-size=128000"], 128000),
        (["-c", "128000"], 128000),
        (["-c=128000"], 128000),
        (["-c", "4096", "--ctx-size", "128000"], 128000),
    ],
)
def test_parse_ctx_override(args, expected):
    assert parse_ctx_override(args) == expected


@pytest.mark.parametrize(
    "args",
    [
        ["--ctx-size"],
        ["--ctx-size", "--top-k"],
        ["--ctx-size", "abc"],
        ["--ctx-size=abc"],
        ["-c", "-1"],
    ],
)
def test_parse_ctx_override_rejects_malformed_values(args):
    with pytest.raises(ValueError, match = "ctx-size|'-c'"):
        parse_ctx_override(args)


def test_validate_extra_args_rejects_malformed_ctx_override():
    with pytest.raises(ValueError, match = "ctx-size"):
        validate_extra_args(["--ctx-size", "abc"])


# ── parse_cache_override ─────────────────────────────────────────────


@pytest.mark.parametrize(
    "args,expected",
    [
        (None, None),
        ([], None),
        (["--top-k", "20"], None),
        (["--cache-type-k", "q8_0"], "q8_0"),
        (["-ctk", "q4_0"], "q4_0"),
        (["-ctv", "q4_0"], "q4_0"),
        (["--cache-type-k=q4_0"], "q4_0"),
        (["-ctk", "f16", "-ctk", "q8_0"], "q8_0"),
    ],
)
def test_parse_cache_override(args, expected):
    assert parse_cache_override(args) == expected


@pytest.mark.parametrize(
    "args",
    [
        ["-ctk"],
        ["-ctk", "-c", "4096"],
    ],
)
def test_parse_cache_override_rejects_malformed_values(args):
    with pytest.raises(ValueError, match = "cache-type|'-ctk'"):
        parse_cache_override(args)


def test_resolve_cache_type_kv_uses_override_when_present():
    assert resolve_cache_type_kv(["--cache-type-k", "q8_0"], "f16") == "q8_0"


def test_resolve_cache_type_kv_uses_fallback_without_override():
    assert resolve_cache_type_kv(["--top-k", "20"], "f16") == "f16"


def test_strip_shadowing_flags_boolean_does_not_consume_next_token():
    # `--spec-default` is boolean; drop just the flag, keep the next token.
    out = strip_shadowing_flags(["--spec-default", "ngram-mod"], strip_spec = True)
    assert out == ["ngram-mod"]


def test_strip_shadowing_flags_jinja_boolean_preserves_positional():
    out = strip_shadowing_flags(["--jinja", "trailing-positional"], strip_template = True)
    assert out == ["trailing-positional"]


def test_strip_shadowing_flags_no_jinja_boolean_preserves_positional():
    out = strip_shadowing_flags(["--no-jinja", "trailing-positional"], strip_template = True)
    assert out == ["trailing-positional"]


def test_strip_shadowing_flags_equals_form_drops_only_the_flag():
    out = strip_shadowing_flags(["--ctx-size=4096", "--seed", "-1"], strip_context = True)
    assert out == ["--seed", "-1"]


def test_strip_shadowing_flags_handles_none_input():
    assert strip_shadowing_flags(None) == []


def test_strip_shadowing_flags_handles_empty_input():
    assert strip_shadowing_flags([]) == []


def test_strip_shadowing_flags_defaults_strip_everything():
    # The route's already-loaded comparator calls with no kwargs to
    # detect ANY shadowing flag in stored extras.
    out = strip_shadowing_flags(
        ["-c", "4096", "--cache-type-k", "q8_0", "--spec-default", "--jinja"]
    )
    assert out == []


def test_strip_shadowing_flags_drops_model_draft_with_spec():
    # --model-draft (and aliases) are Studio-managed since the separate
    # MTP drafter support: an inherited copy must not last-wins-override
    # the auto-detected drafter.
    out = strip_shadowing_flags(
        ["--model-draft", "/old/mtp.gguf", "-md", "/old2.gguf", "--top-k", "20"],
        strip_context = False,
        strip_cache = False,
        strip_spec = True,
        strip_template = False,
    )
    assert out == ["--top-k", "20"]


def test_strip_shadowing_flags_keeps_model_draft_without_spec():
    out = strip_shadowing_flags(
        ["--model-draft", "/custom/mtp.gguf"],
        strip_context = True,
        strip_cache = False,
        strip_spec = False,
        strip_template = False,
    )
    assert out == ["--model-draft", "/custom/mtp.gguf"]