import json
import os
import re
import shlex
import stat
import subprocess
import sys
from pathlib import Path
from typing import Final
from urllib.parse import parse_qs, urlencode, urlsplit

import pytest

from litellm.constants import PROMETHEUS_ADMITTED_SERIES_FILE_PREFIX

REPO_ROOT: Final = Path(__file__).resolve().parents[2]
DOCKER_ENTRYPOINT: Final = REPO_ROOT / "litellm/deploy/docker/docker-entrypoint.sh"
DOCKERFILE: Final = REPO_ROOT / "litellm/deploy/docker/Dockerfile"
TERRAFORM_MODULES: Final = {
    "aws": REPO_ROOT / "litellm" / "deploy" / "terraform" / "litellm" / "aws",
    "gcp": REPO_ROOT / "litellm" / "deploy" / "terraform" / "litellm" / "gcp",
}

IMAGE_ENTRYPOINT_PATH: Final = "/app/docker-entrypoint.sh"
STUBBED_EXECUTABLES: Final = ("ddtrace-run", "uvicorn", "python", "litellm", "nginx")

TRUTHY_USE_DDTRACE: Final = ("true", "True", "TRUE", "tRuE")
FALSY_USE_DDTRACE: Final = (None, "", "false", "False", "1", "yes", "on", "truex")

PYTHONPATH_SENTINEL: Final = "/lit-entrypoint-sentinel:/app"

_STUB_TEMPLATE: Final = """#!/bin/sh
{{
  echo "exec={name}"
  echo "args=$*"
  echo "DD_TRACE_OPENAI_ENABLED=${{DD_TRACE_OPENAI_ENABLED-<unset>}}"
  echo "PYTHONPATH=${{PYTHONPATH-<unset>}}"
}} >> "$RECORD"
"""

_ENTRYPOINT_RE: Final = re.compile(r"^ENTRYPOINT\s+(\[.*\])\s*$", re.MULTILINE)
_CMD_RE: Final = re.compile(r"^CMD\s+(\[.*\])\s*$", re.MULTILINE)
_TF_ENTRYPOINT_LOCAL_RE: Final = re.compile(r'^\s*image_entrypoint\s*=\s*"([^"]*)"\s*$', re.MULTILINE)
_TF_WORD: Final = r'(?:"[^"\s]*"|tostring\(var\.\w+\)|\$\{var\.\w+\}|[\w.-]+)'
_TF_ENTRYPOINT_CALL_RE: Final = re.compile(rf"exec \$\{{local\.image_entrypoint\}}((?: {_TF_WORD})+)")
_TF_COMPONENT_LIST_RE: Final = re.compile(
    rf"(?:command|args|gateway|backend|ui|migrations|metrics|collector)\s*=\s*"
    rf'\[\s*("(?:proxy|gateway|backend|ui|migrations|metrics|collector)"(?:\s*,\s*{_TF_WORD})*)\s*\]'
)
_TF_VAR_RE: Final = re.compile(r"tostring\(var\.(\w+)\)|\$\{var\.(\w+)\}")

TERRAFORM_VAR_STUBS: Final = {"gateway_num_workers": "2", "gateway_metrics_port": "9464"}
TERRAFORM_WORKLOADS: Final = frozenset({"gateway", "backend", "ui", "migrations", "metrics", "collector"})
COMPONENT_BINARIES: Final = {
    "gateway": "python",
    "backend": "uvicorn",
    "ui": "nginx",
    "migrations": "python",
    "metrics": "python",
    "collector": "python",
}


def _write_stubs(bin_dir: Path, names: tuple[str, ...]) -> None:
    for name in names:
        stub = bin_dir / name
        stub.write_text(_STUB_TEMPLATE.format(name=name))
        stub.chmod(0o755)


def _entrypoint_env(bin_dir: Path, record: Path, overrides: dict[str, str | None]) -> dict[str, str]:
    cleared = ("USE_DDTRACE", "DD_TRACE_OPENAI_ENABLED", "LITELLM_COMPONENT", "NUM_WORKERS", *overrides)
    return {
        **{k: v for k, v in os.environ.items() if k not in cleared},
        **{k: v for k, v in overrides.items() if v is not None},
        "PATH": f"{bin_dir}{os.pathsep}{os.environ['PATH']}",
        "RECORD": str(record),
        "PYTHONPATH": PYTHONPATH_SENTINEL,
    }


def _run_entrypoint(
    argv: tuple[str, ...],
    tmp_path: Path,
    use_ddtrace: str | None = None,
    *,
    entrypoint: Path = DOCKER_ENTRYPOINT,
    **overrides: str | None,
) -> tuple[str, ...]:
    bin_dir = tmp_path / "bin"
    bin_dir.mkdir(parents=True)
    _write_stubs(bin_dir, STUBBED_EXECUTABLES)
    record = tmp_path / "record.txt"
    env = _entrypoint_env(bin_dir, record, {"USE_DDTRACE": use_ddtrace, **overrides})

    result = subprocess.run(["sh", str(entrypoint), *argv], env=env, capture_output=True, text=True, check=False)
    assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}"
    return tuple(record.read_text().splitlines()) if record.exists() else ()


def _entrypoint_with_runtime_config(tmp_path: Path) -> tuple[Path, Path]:
    entrypoint: Final = tmp_path / "docker-entrypoint.sh"
    entrypoint.write_text(DOCKER_ENTRYPOINT.read_text())
    hook: Final = tmp_path / "ui" / "docker-entrypoint.d" / "40-litellm-runtime-config.sh"
    hook.parent.mkdir(parents=True)
    runtime_dir: Final = tmp_path / "runtime"
    runtime_dir.mkdir()
    hook.write_text(
        (REPO_ROOT / "litellm" / "ui" / "docker-entrypoint.d" / hook.name)
        .read_text()
        .replace("runtime_dir=/tmp/litellm-ui", f"runtime_dir={shlex.quote(str(runtime_dir))}")
    )
    return entrypoint, runtime_dir / "runtime-config.json"


@pytest.mark.parametrize("backend_url", ("https://backend.example.com///", "http://backend.example.com", None, "", "invalid"))
def test_ui_runtime_config_is_refreshed_before_nginx_starts(tmp_path: Path, backend_url: str | None) -> None:
    entrypoint, runtime_file = _entrypoint_with_runtime_config(tmp_path)
    runtime_file.write_text('{"backend_url":"https://stale.example.com"}')
    recorded: Final = _run_entrypoint(("ui",), tmp_path, entrypoint=entrypoint, LITELLM_BACKEND_URL=backend_url)
    assert recorded[0] == "exec=nginx"
    assert (json.loads(runtime_file.read_text()) if runtime_file.exists() else None) == (
        {"backend_url": backend_url.rstrip("/")} if backend_url and backend_url.startswith(("https://", "http://")) else None
    )


def test_backend_leaves_ui_runtime_config_unchanged(tmp_path: Path) -> None:
    entrypoint, runtime_file = _entrypoint_with_runtime_config(tmp_path)
    runtime_file.write_text('{"backend_url":"https://stale.example.com"}')
    recorded: Final = _run_entrypoint(
        ("backend",), tmp_path, entrypoint=entrypoint, LITELLM_BACKEND_URL="https://new.example.com"
    )
    assert recorded[0] == "exec=uvicorn"
    assert json.loads(runtime_file.read_text()) == {"backend_url": "https://stale.example.com"}


def _tf_word(word: str) -> str:
    var: Final = _TF_VAR_RE.fullmatch(word)
    if var is not None:
        return TERRAFORM_VAR_STUBS[var.group(1) or var.group(2)]
    return word.strip('"')


def _terraform_component_argvs(module_dir: Path) -> tuple[tuple[str, tuple[str, ...]], ...]:
    return tuple(
        (tf.name, tuple(_tf_word(word) for word in words))
        for tf in sorted(module_dir.glob("*.tf"))
        for words in (
            *(m.group(1).split() for m in _TF_ENTRYPOINT_CALL_RE.finditer(tf.read_text())),
            *(re.split(r"\s*,\s*", m.group(1)) for m in _TF_COMPONENT_LIST_RE.finditer(tf.read_text())),
        )
    )  # comprehension-ok: the two regex scans over one file read are clearer inline than as a generator helper


def _entrypoint_argv(dockerfile: Path) -> tuple[str, ...]:
    matches = _ENTRYPOINT_RE.findall(dockerfile.read_text())
    assert matches, f"no exec-form ENTRYPOINT found in {dockerfile}"
    return tuple(str(part) for part in json.loads(matches[-1]))


@pytest.mark.parametrize(
    "argv, expected_exec, expected_args",
    [
        (("proxy", "--config", "/app/config.yaml"), "exec=litellm", "args=--config /app/config.yaml"),
        (("gateway",), "exec=python", "args=-m gateway.launch --workers 1 --host 0.0.0.0 --port 4000"),
        (
            ("gateway", "--port", "8080"),
            "exec=python",
            "args=-m gateway.launch --workers 1 --host 0.0.0.0 --port 4000 --port 8080",
        ),
        (("backend",), "exec=uvicorn", "args=litellm.runtime.backend.main:app --host 0.0.0.0 --port 4001"),
        (("backend", "--port", "9001"), "exec=uvicorn", "args=litellm.runtime.backend.main:app --host 0.0.0.0 --port 4001 --port 9001"),
        (("ui",), "exec=nginx", "args=-g daemon off;"),
        (("migrations",), "exec=python", "args=-m litellm.runtime.migrations.run"),
        (("metrics", "--port", "9090"), "exec=python", "args=-m litellm.proxy.prometheus_metrics_server --port 9090"),
        (("collector",), "exec=python", "args=-m litellm.proxy.collector"),
    ],
)
def test_the_first_argument_selects_the_component(
    argv: tuple[str, ...], expected_exec: str, expected_args: str, tmp_path: Path
) -> None:
    recorded = _run_entrypoint(argv, tmp_path)

    assert recorded[:2] == (expected_exec, expected_args)


@pytest.mark.parametrize(
    "argv, expected",
    [
        ((), ("exec=litellm", "args=")),
        (("--port", "4000"), ("exec=litellm", "args=--port 4000")),
        (
            ("--config", "/app/config.yaml", "--detailed_debug"),
            ("exec=litellm", "args=--config /app/config.yaml --detailed_debug"),
        ),
    ],
)
def test_flags_alone_still_run_the_monolithic_proxy(
    argv: tuple[str, ...], expected: tuple[str, str], tmp_path: Path
) -> None:
    assert _run_entrypoint(argv, tmp_path)[:2] == expected


def test_the_dockerfile_leaves_the_command_empty_so_the_env_var_can_pick_the_component(tmp_path: Path) -> None:
    dockerfile_text = DOCKERFILE.read_text()
    assert _entrypoint_argv(DOCKERFILE) == (IMAGE_ENTRYPOINT_PATH,), "the image ENTRYPOINT must be the bare dispatcher"
    assert not _CMD_RE.search(dockerfile_text), "the Dockerfile must not set a CMD"

    assert _run_entrypoint((), tmp_path)[:2] == ("exec=litellm", "args=")


@pytest.mark.parametrize(
    "component, argv, expected_exec, expected_args",
    [
        ("backend", (), "exec=uvicorn", "args=litellm.runtime.backend.main:app --host 0.0.0.0 --port 4001"),
        (
            "gateway",
            ("--port", "4100"),
            "exec=python",
            "args=-m gateway.launch --workers 1 --host 0.0.0.0 --port 4000 --port 4100",
        ),
        ("ui", (), "exec=nginx", "args=-g daemon off;"),
    ],
)
def test_litellm_component_env_selects_the_component_when_the_command_has_none(
    component: str, argv: tuple[str, ...], expected_exec: str, expected_args: str, tmp_path: Path
) -> None:
    recorded = _run_entrypoint(argv, tmp_path, LITELLM_COMPONENT=component)

    assert recorded[:2] == (expected_exec, expected_args)


def test_an_explicit_component_argument_beats_the_env_var(tmp_path: Path) -> None:
    recorded = _run_entrypoint(("backend",), tmp_path, LITELLM_COMPONENT="gateway")

    assert recorded[0] == "exec=uvicorn"


def test_the_gateway_honours_num_workers(tmp_path: Path) -> None:
    recorded = _run_entrypoint(("gateway",), tmp_path, NUM_WORKERS="4")

    assert recorded[1] == "args=-m gateway.launch --workers 4 --host 0.0.0.0 --port 4000"


def test_an_unknown_first_word_is_run_as_the_command(tmp_path: Path) -> None:
    bin_dir = tmp_path / "bin"
    bin_dir.mkdir()
    _write_stubs(bin_dir, ("some-tool",))
    record = tmp_path / "record.txt"
    env = _entrypoint_env(bin_dir, record, {})

    result = subprocess.run(
        ["sh", str(DOCKER_ENTRYPOINT), "some-tool", "--flag"], env=env, capture_output=True, text=True, check=False
    )

    assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}"
    assert record.read_text().splitlines()[:2] == ["exec=some-tool", "args=--flag"]


@pytest.mark.parametrize(
    "use_ddtrace, traced",
    [*((v, True) for v in TRUTHY_USE_DDTRACE), *((v, False) for v in FALSY_USE_DDTRACE)],
)
def test_ddtrace_wraps_a_raw_command_like_a_component(use_ddtrace: str | None, traced: bool, tmp_path: Path) -> None:
    bin_dir = tmp_path / "bin"
    bin_dir.mkdir()
    _write_stubs(bin_dir, ("some-tool", "ddtrace-run"))
    record = tmp_path / "record.txt"
    env = _entrypoint_env(bin_dir, record, {"USE_DDTRACE": use_ddtrace})

    result = subprocess.run(
        ["sh", str(DOCKER_ENTRYPOINT), "some-tool", "--flag"], env=env, capture_output=True, text=True, check=False
    )

    assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}"
    assert record.read_text().splitlines()[:2] == (
        ["exec=ddtrace-run", "args=some-tool --flag"] if traced else ["exec=some-tool", "args=--flag"]
    )


def test_an_unknown_litellm_component_fails_fast_instead_of_guessing(tmp_path: Path) -> None:
    bin_dir = tmp_path / "bin"
    bin_dir.mkdir()
    _write_stubs(bin_dir, STUBBED_EXECUTABLES)
    record = tmp_path / "record.txt"
    env = _entrypoint_env(bin_dir, record, {"LITELLM_COMPONENT": "gatway"})

    result = subprocess.run(["sh", str(DOCKER_ENTRYPOINT)], env=env, capture_output=True, text=True, check=False)

    assert result.returncode == 64
    assert "gatway" in result.stderr
    assert not record.exists(), "nothing may be exec'd when the component is unknown"


def test_admin_agent_flag_runs_the_bundled_liteadmin_agent_not_the_proxy(tmp_path: Path) -> None:
    bin_dir = tmp_path / "bin"
    bin_dir.mkdir()
    _write_stubs(bin_dir, STUBBED_EXECUTABLES)
    record = tmp_path / "record.txt"
    env = _entrypoint_env(bin_dir, record, {})

    result = subprocess.run(
        ["sh", str(DOCKER_ENTRYPOINT), "--admin-agent", "--flag"], env=env, capture_output=True, text=True, check=False
    )

    assert "/opt/liteadmin/bin/litellm-admin-agent" in result.stderr
    assert result.returncode == 127
    assert not record.exists(), "--admin-agent must not fall through to the proxy"


def test_ddtrace_enabled_wraps_the_command_and_disables_the_openai_integration(tmp_path: Path) -> None:
    recorded = _run_entrypoint(("gateway",), tmp_path, use_ddtrace="true", NUM_WORKERS="2")

    assert recorded == (
        "exec=ddtrace-run",
        "args=python -m gateway.launch --workers 2 --host 0.0.0.0 --port 4000",
        "DD_TRACE_OPENAI_ENABLED=False",
        f"PYTHONPATH={PYTHONPATH_SENTINEL}",
    )


def test_ddtrace_disabled_execs_the_command_directly(tmp_path: Path) -> None:
    recorded = _run_entrypoint(("backend", "--workers", "2"), tmp_path)

    assert recorded == (
        "exec=uvicorn",
        "args=litellm.runtime.backend.main:app --host 0.0.0.0 --port 4001 --workers 2",
        "DD_TRACE_OPENAI_ENABLED=<unset>",
        f"PYTHONPATH={PYTHONPATH_SENTINEL}",
    )


@pytest.mark.parametrize(
    "use_ddtrace, traced",
    [*((v, True) for v in TRUTHY_USE_DDTRACE), *((v, False) for v in FALSY_USE_DDTRACE)],
)
def test_ddtrace_gating_matches_get_secret_bool_for_every_component(
    use_ddtrace: str | None, traced: bool, tmp_path: Path
) -> None:
    component = _run_entrypoint(("gateway",), tmp_path / "component", use_ddtrace=use_ddtrace)
    monolith = _run_entrypoint(("--port", "4000"), tmp_path / "monolith", use_ddtrace=use_ddtrace)

    expected_openai = "DD_TRACE_OPENAI_ENABLED=False" if traced else "DD_TRACE_OPENAI_ENABLED=<unset>"
    assert component[0] == ("exec=ddtrace-run" if traced else "exec=python")
    assert component[2] == expected_openai
    assert monolith[0] == ("exec=ddtrace-run" if traced else "exec=litellm")
    assert monolith[1] == ("args=litellm --port 4000" if traced else "args=--port 4000")
    assert monolith[2] == expected_openai


def test_wipes_the_prometheus_multiproc_dir_before_uvicorn_forks(tmp_path: Path) -> None:
    multiproc_dir = tmp_path / "multiproc"
    multiproc_dir.mkdir()
    (multiproc_dir / "gauge_livesum_7.db").write_bytes(b"stale")
    (multiproc_dir / "counter_7.db").write_bytes(b"stale")
    (multiproc_dir / f"{PROMETHEUS_ADMITTED_SERIES_FILE_PREFIX}litellm_requests_metric").write_bytes(b"stale")
    (multiproc_dir / "keep.txt").write_text("not a sample")

    recorded = _run_entrypoint(("gateway",), tmp_path, PROMETHEUS_MULTIPROC_DIR=str(multiproc_dir))

    assert sorted(p.name for p in multiproc_dir.iterdir()) == ["keep.txt"]
    assert recorded[0] == "exec=python"


@pytest.mark.parametrize(
    "argv",
    [
        ("python", "-m", "litellm.proxy.prometheus_metrics_server"),
        ("python", "-m", "litellm.proxy.collector"),
        ("sh", "-c", "exec python -m litellm.proxy.collector --port 4010"),
    ],
)
def test_a_raw_metrics_reader_keeps_the_workers_samples(argv: tuple[str, ...], tmp_path: Path) -> None:
    multiproc_dir = tmp_path / "multiproc"
    multiproc_dir.mkdir()
    (multiproc_dir / "counter_7.db").write_bytes(b"live")
    (multiproc_dir / f"{PROMETHEUS_ADMITTED_SERIES_FILE_PREFIX}litellm_requests_metric").write_bytes(b"live")
    bin_dir = tmp_path / "bin"
    bin_dir.mkdir()
    _write_stubs(bin_dir, ("python",))
    record = tmp_path / "record.txt"
    env = _entrypoint_env(bin_dir, record, {"PROMETHEUS_MULTIPROC_DIR": str(multiproc_dir)})

    result = subprocess.run(
        ["sh", str(DOCKER_ENTRYPOINT), *argv],
        env=env,
        capture_output=True,
        text=True,
        check=False,
    )

    assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}"
    assert sorted(p.name for p in multiproc_dir.iterdir()) == [
        "counter_7.db",
        f"{PROMETHEUS_ADMITTED_SERIES_FILE_PREFIX}litellm_requests_metric",
    ]
    assert record.read_text().splitlines()[0] == "exec=python"


@pytest.mark.parametrize(
    "argv",
    [
        ("python", "-m", "gateway.launch"),
        ("litellm", "--config", "/app/config.yaml"),
        ("uvicorn", "litellm.runtime.backend.main:app"),
    ],
)
def test_a_raw_writer_starts_with_an_empty_prometheus_multiproc_dir(argv: tuple[str, ...], tmp_path: Path) -> None:
    multiproc_dir = tmp_path / "multiproc"
    multiproc_dir.mkdir()
    (multiproc_dir / "counter_7.db").write_bytes(b"stale")
    (multiproc_dir / "gauge_livesum_7.db").write_bytes(b"stale")
    (multiproc_dir / f"{PROMETHEUS_ADMITTED_SERIES_FILE_PREFIX}litellm_requests_metric").write_bytes(b"stale")
    (multiproc_dir / "keep.txt").write_text("not a sample")
    bin_dir = tmp_path / "bin"
    bin_dir.mkdir()
    _write_stubs(bin_dir, STUBBED_EXECUTABLES)
    record = tmp_path / "record.txt"
    env = _entrypoint_env(bin_dir, record, {"PROMETHEUS_MULTIPROC_DIR": str(multiproc_dir)})

    result = subprocess.run(
        ["sh", str(DOCKER_ENTRYPOINT), *argv],
        env=env,
        capture_output=True,
        text=True,
        check=False,
    )

    assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}"
    assert sorted(p.name for p in multiproc_dir.iterdir()) == ["keep.txt"]
    assert record.read_text().splitlines()[0] == f"exec={argv[0]}"


def test_a_raw_command_still_gets_the_prometheus_multiproc_dir_created(tmp_path: Path) -> None:
    multiproc_dir = tmp_path / "not-yet" / "multiproc"
    bin_dir = tmp_path / "bin"
    bin_dir.mkdir()
    _write_stubs(bin_dir, ("python",))
    record = tmp_path / "record.txt"
    env = _entrypoint_env(bin_dir, record, {"PROMETHEUS_MULTIPROC_DIR": str(multiproc_dir)})

    result = subprocess.run(
        ["sh", str(DOCKER_ENTRYPOINT), "python", "-m", "gateway.launch"],
        env=env,
        capture_output=True,
        text=True,
        check=False,
    )

    assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}"
    assert multiproc_dir.is_dir()
    assert record.read_text().splitlines()[0] == "exec=python"


@pytest.mark.parametrize("component", ["metrics", "collector"])
def test_readers_of_the_prometheus_multiproc_dir_keep_the_workers_samples(component: str, tmp_path: Path) -> None:
    multiproc_dir = tmp_path / "multiproc"
    multiproc_dir.mkdir()
    (multiproc_dir / "counter_12.db").write_bytes(b"live")
    (multiproc_dir / "histogram_13.db").write_bytes(b"live")
    (multiproc_dir / f"{PROMETHEUS_ADMITTED_SERIES_FILE_PREFIX}litellm_requests_metric").write_bytes(b"live")

    recorded = _run_entrypoint((component,), tmp_path, PROMETHEUS_MULTIPROC_DIR=str(multiproc_dir))

    assert sorted(p.name for p in multiproc_dir.iterdir()) == [
        "counter_12.db",
        "histogram_13.db",
        f"{PROMETHEUS_ADMITTED_SERIES_FILE_PREFIX}litellm_requests_metric",
    ]
    assert recorded[0] == "exec=python"


def test_creates_a_missing_prometheus_multiproc_dir(tmp_path: Path) -> None:
    missing = tmp_path / "multiproc"

    _run_entrypoint(("backend",), tmp_path, PROMETHEUS_MULTIPROC_DIR=str(missing))

    assert missing.is_dir()


@pytest.mark.parametrize("module", TERRAFORM_MODULES, ids=str)
def test_terraform_hands_every_workload_to_the_image_entrypoint(module: str) -> None:
    argvs: Final = _terraform_component_argvs(TERRAFORM_MODULES[module])

    assert {argv[0] for _, argv in argvs} == TERRAFORM_WORKLOADS, f"{module} workloads: {argvs}"
    for tf in TERRAFORM_MODULES[module].glob("*.tf"):
        for entrypoint_path in _TF_ENTRYPOINT_LOCAL_RE.findall(tf.read_text()):
            assert entrypoint_path == _entrypoint_argv(DOCKERFILE)[0], f"{tf} execs a path the image does not ship"


@pytest.mark.parametrize("module", TERRAFORM_MODULES, ids=str)
@pytest.mark.parametrize("use_ddtrace", ["true", None])
def test_the_entrypoint_accepts_every_terraform_argv(module: str, use_ddtrace: str | None, tmp_path: Path) -> None:
    for index, (tf, argv) in enumerate(_terraform_component_argvs(TERRAFORM_MODULES[module])):
        component, *extra = argv
        recorded: Final = _run_entrypoint(argv, tmp_path / str(index), use_ddtrace=use_ddtrace)

        expected_exec: Final = "ddtrace-run" if use_ddtrace and component != "ui" else COMPONENT_BINARIES[component]
        assert recorded[0] == f"exec={expected_exec}", f"{module}/{tf} {argv}: {recorded}"
        assert recorded[1].endswith(" ".join(extra)), f"{module}/{tf} {argv} drops its flags: {recorded[1]}"
        if component == "gateway":
            assert extra == ["--workers", TERRAFORM_VAR_STUBS["gateway_num_workers"]], f"{module}/{tf}: {argv}"
        if component == "metrics":
            assert extra == ["--port", TERRAFORM_VAR_STUBS["gateway_metrics_port"]], f"{module}/{tf}: {argv}"


def test_help_prints_the_header_comment_block_without_shell_source(tmp_path: Path) -> None:
    bin_dir = tmp_path / "bin"
    bin_dir.mkdir(parents=True)
    _write_stubs(bin_dir, STUBBED_EXECUTABLES)
    env = _entrypoint_env(bin_dir, tmp_path / "record.txt", {})

    result = subprocess.run(
        ["sh", str(DOCKER_ENTRYPOINT), "help"], env=env, capture_output=True, text=True, check=False
    )

    assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}"
    assert "--admin-agent" in result.stderr
    assert "collector" in result.stderr
    assert "set -eu" not in result.stderr
    assert result.stderr.rstrip().endswith("other raw commands clear stale samples.")


def test_entrypoint_script_is_executable() -> None:
    mode = DOCKER_ENTRYPOINT.stat().st_mode
    assert mode & stat.S_IXUSR, "entrypoint must be committed executable to run as the image ENTRYPOINT"
    assert mode & stat.S_IXOTH, "entrypoint must be executable by the unprivileged `nonroot` user"


def test_entrypoint_script_has_no_carriage_returns() -> None:
    assert b"\r" not in DOCKER_ENTRYPOINT.read_bytes()


def _run_proxy_with_discrete_database(
    tmp_path: Path, extra_env: dict[str, str | None], extra_pythonpath: Path | None = None
) -> list[str]:
    bin_dir = tmp_path / "bin"
    bin_dir.mkdir()
    (bin_dir / "litellm").write_text(
        '#!/bin/sh\n{ echo "args=$*"; echo "DATABASE_URL=$DATABASE_URL"; '
        'echo "DATABASE_URL_READ_REPLICA=${DATABASE_URL_READ_REPLICA-<unset>}"; } >> "$RECORD"\n'
    )
    (bin_dir / "litellm").chmod(0o755)
    (bin_dir / "python").write_text(f'#!/bin/sh\nexec "{sys.executable}" "$@"\n')
    (bin_dir / "python").chmod(0o755)
    record = tmp_path / "record.txt"
    env = _entrypoint_env(
        bin_dir,
        record,
        {
            "DATABASE_URL": None,
            "DATABASE_URL_READ_REPLICA": None,
            "DIRECT_URL": None,
            "DATABASE_MAX_IDLE_CONNECTION_LIFETIME": None,
            "USE_AWS_KMS": None,
            "DATABASE_HOST": "writer.db",
            "DATABASE_PORT": "6543",
            "DATABASE_USER": "app",
            "DATABASE_PASSWORD": "p@ss",
            "DATABASE_NAME": "litellm",
            "DATABASE_HOST_READ_REPLICA": "reader.db",
            "DATABASE_PORT_READ_REPLICA": "6544",
            **extra_env,
        },
    )
    env["PYTHONPATH"] = os.pathsep.join(
        [str(extra_pythonpath), str(REPO_ROOT)] if extra_pythonpath else [str(REPO_ROOT)]
    )

    result = subprocess.run(
        ["sh", str(DOCKER_ENTRYPOINT), "proxy", "--port", "4000"],
        env=env,
        capture_output=True,
        text=True,
        check=False,
    )

    assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}"
    return record.read_text().splitlines()


def test_proxy_assembles_database_urls_from_discrete_variables(tmp_path: Path) -> None:
    lines = _run_proxy_with_discrete_database(tmp_path, {})

    assert lines[0] == "args=--port 4000"
    assert lines[1].startswith("DATABASE_URL=postgresql://app:p%40ss@writer.db:6543/litellm")
    assert lines[2].startswith("DATABASE_URL_READ_REPLICA=postgresql://app:p%40ss@reader.db:6544/litellm")
    assert "max_idle_connection_lifetime" not in parse_qs(urlsplit(lines[1]).query)
    assert "max_idle_connection_lifetime" not in parse_qs(urlsplit(lines[2]).query)


def test_proxy_decrypts_a_kms_database_password_before_assembling(tmp_path: Path) -> None:
    site_dir = tmp_path / "site"
    site_dir.mkdir()
    (site_dir / "sitecustomize.py").write_text(
        "import os\n"
        "import litellm.secret_managers.aws_secret_manager as kms\n"
        "kms.decrypt_env_var = lambda: {'DATABASE_PASSWORD': 'p@ss'} "
        "if os.environ.get('DATABASE_PASSWORD') == 'aws_kms/ciphertext' else {}\n"
    )
    lines = _run_proxy_with_discrete_database(
        tmp_path,
        {"USE_AWS_KMS": "True", "DATABASE_PASSWORD": "aws_kms/ciphertext"},
        extra_pythonpath=site_dir,
    )

    assert lines[1].startswith("DATABASE_URL=postgresql://app:p%40ss@writer.db:6543/litellm")
    assert lines[2].startswith("DATABASE_URL_READ_REPLICA=postgresql://app:p%40ss@reader.db:6544/litellm")


def test_proxy_pins_an_env_configured_idle_lifetime(tmp_path: Path) -> None:
    lines = _run_proxy_with_discrete_database(tmp_path, {"DATABASE_MAX_IDLE_CONNECTION_LIFETIME": "120"})

    assert "max_idle_connection_lifetime=120" in lines[1]
    assert "max_idle_connection_lifetime=120" in lines[2]


def test_proxy_keeps_a_reader_url_pinned_idle_lifetime(tmp_path: Path) -> None:
    lines: Final = _run_proxy_with_discrete_database(
        tmp_path, {"DATABASE_URL_READ_REPLICA": "postgresql://app:pass@reader.db/litellm?max_idle_connection_lifetime=300"}
    )
    assert "max_idle_connection_lifetime" not in parse_qs(urlsplit(lines[1]).query)
    assert "max_idle_connection_lifetime=300" in lines[2]
    assert "max_idle_connection_lifetime=60" not in lines[2]


@pytest.mark.parametrize("pinned_writer", (False, True))
def test_proxy_reader_uses_prisma_tls_and_writer_pool_settings(tmp_path: Path, pinned_writer: bool) -> None:
    root_cert: Final = str(tmp_path / "absent-ca.pem")
    writer_query: Final = urlencode(
        {
            "sslmode": "verify-full",
            "sslrootcert": root_cert,
            "connection_limit": "7",
            "socket_timeout": "9",
            "max_idle_connection_lifetime": "300",
        }
    )
    lines: Final = _run_proxy_with_discrete_database(
        tmp_path,
        {
            "DATABASE_URL": f"postgresql://app:pass@writer.db/litellm?{writer_query}" if pinned_writer else None,
            "DATABASE_HOST": None if pinned_writer else "writer.db",
            "DATABASE_USER_READ_REPLICA": "reader",
            "DATABASE_PASSWORD_READ_REPLICA": "rpass",
            "DATABASE_NAME_READ_REPLICA": "replica",
            "DATABASE_SSLMODE": "verify-full",
            "DATABASE_SSLROOTCERT": root_cert,
        },
    )
    reader: Final = urlsplit(lines[2].removeprefix("DATABASE_URL_READ_REPLICA="))
    assert (reader.hostname, reader.username, reader.path) == ("reader.db", "reader", "/replica")
    assert parse_qs(reader.query) == {
        "sslmode": ["require"],
        "sslcert": [root_cert],
        "sslaccept": ["strict"],
        **(
            {"connection_limit": ["7"], "socket_timeout": ["9"], "max_idle_connection_lifetime": ["300"]}
            if pinned_writer
            else {}
        ),
    }


def test_proxy_keeps_an_operator_database_url(tmp_path: Path) -> None:
    record = _run_entrypoint(
        ("proxy",),
        tmp_path,
        DATABASE_URL="postgresql://pinned/db",
        DATABASE_HOST="writer.db",
    )

    assert record[0] == "exec=litellm"