tests/test_jtp3_wrapper.py

from __future__ import annotations

import os
import subprocess
import sys
from pathlib import Path

import pytest

import kneiff.utils.image.tag_jtp3 as tag_jtp3


UPSTREAM_TAGS = r"from side, artist \(name\)"


def _write_current_hydra_snapshot(snapshot: Path) -> None:
    snapshot.mkdir()
    (snapshot / "inference.py").write_text(
        'group.add_argument("-b", "--batch-size", type=int, default=1)\n',
        encoding="utf-8",
    )


def _patch_jtp3_subprocess(
    tmp_path: Path,
    monkeypatch,
) -> dict[str, object]:
    snapshot = tmp_path / "snapshot"
    calls: dict[str, object] = {}

    def fake_snapshot_download(**kwargs: object) -> str:
        calls["snapshot_kwargs"] = kwargs
        return str(snapshot)

    def fake_run(
        command: list[str],
        *,
        cwd: Path,
        env: dict[str, str],
        check: bool,
        text: bool,
    ) -> subprocess.CompletedProcess[str]:
        calls["run"] = (command, cwd, env, check, text)
        for value in command:
            path = Path(value)
            if path.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp"}:
                path.with_suffix(".txt").write_text(
                    f"{UPSTREAM_TAGS}\n", encoding="utf-8"
                )
        return subprocess.CompletedProcess(command, 0)

    monkeypatch.setattr(tag_jtp3, "snapshot_download", fake_snapshot_download)
    monkeypatch.setattr(tag_jtp3.subprocess, "run", fake_run)
    return calls


def test_build_jtp3_command_uses_stdout_csv_and_snapshot_cwd(tmp_path: Path) -> None:
    snapshot = tmp_path / "snapshot"
    config = tag_jtp3.Jtp3RunConfig(
        paths=[Path("images")],
        recursive=True,
        threshold=0.35,
        device="cpu",
        batch_size=4,
        workers=2,
        seqlen=2048,
        prefix="knf-",
        csv_stdout=True,
    )

    command = tag_jtp3.build_jtp3_command(snapshot, config)

    assert command[:2] == [sys.executable, str(snapshot / "inference.py")]
    assert "--output" in command
    assert "-" in command
    assert "--recursive" in command
    assert ["--device", "cpu"] == command[
        command.index("--device") : command.index("--device") + 2
    ]


def test_build_jtp3_command_uses_current_hydra_cli_shape(tmp_path: Path) -> None:
    snapshot = tmp_path / "snapshot"
    _write_current_hydra_snapshot(snapshot)
    config = tag_jtp3.Jtp3RunConfig(
        paths=[Path("image.png")],
        recursive=True,
        device="cpu",
        batch_size=4,
        workers=2,
        seqlen=2048,
        prefix="knf-",
        csv_stdout=True,
    )

    command = tag_jtp3.build_jtp3_command(snapshot, config)

    assert command[:4] == [
        sys.executable,
        "-m",
        "kneiff.utils.image.jtp3_upstream_runner",
        str(snapshot / "inference.py"),
    ]
    assert "--threshold" not in command
    assert "--batch" not in command
    assert ["--model", tag_jtp3.CURRENT_HYDRA_JTP3_MODEL] == command[
        command.index("--model") : command.index("--model") + 2
    ]
    assert ["--batch-size", "4"] == command[
        command.index("--batch-size") : command.index("--batch-size") + 2
    ]
    assert ["--seqlen", "2048"] == command[
        command.index("--seqlen") : command.index("--seqlen") + 2
    ]
    assert ["--workers", "2"] == command[
        command.index("--workers") : command.index("--workers") + 2
    ]
    assert ["--prefix", "knf-"] == command[
        command.index("--prefix") : command.index("--prefix") + 2
    ]
    assert ["--device", "cpu"] == command[
        command.index("--device") : command.index("--device") + 2
    ]
    assert "--recursive" in command
    assert "--output" in command
    assert "-" in command
    assert command[-1] == "image.png"


def test_build_jtp3_command_rejects_legacy_threshold_for_current_hydra(
    tmp_path: Path,
) -> None:
    snapshot = tmp_path / "snapshot"
    _write_current_hydra_snapshot(snapshot)
    config = tag_jtp3.Jtp3RunConfig(
        paths=[Path("image.png")],
        threshold=0.35,
    )

    with pytest.raises(ValueError, match="--threshold"):
        tag_jtp3.build_jtp3_command(snapshot, config)


def test_build_jtp3_environment_prepends_snapshot_to_pythonpath(
    tmp_path: Path,
    monkeypatch,
) -> None:
    snapshot = tmp_path / "snapshot"
    monkeypatch.setenv("PYTHONPATH", "existing")

    env = tag_jtp3.build_jtp3_environment(snapshot)

    assert env["PYTHONPATH"] == os.pathsep.join([str(snapshot), "existing"])


def test_parse_jtp3_sidecar_tags_ignores_empty_tokens() -> None:
    assert tag_jtp3.parse_jtp3_sidecar_tags(" alpha, beta ,, gamma\n") == [
        "alpha",
        "beta",
        "gamma",
    ]


def test_format_jtp3_tags_defaults_to_e621_whitespace_output() -> None:
    assert (
        tag_jtp3.format_jtp3_tags(
            [r"from side", r"artist \(name\)"],
            comma_separated=False,
        )
        == "from_side artist_(name)"
    )


def test_format_jtp3_tags_can_keep_legacy_commas() -> None:
    assert (
        tag_jtp3.format_jtp3_tags(
            [r"from side", r"artist \(name\)"],
            comma_separated=True,
        )
        == r"from side, artist \(name\)"
    )


def test_run_jtp3_defaults_to_stdout_without_touching_original_sidecar(
    tmp_path: Path,
    monkeypatch,
    capsys,
) -> None:
    image_path = tmp_path / "image.png"
    image_path.write_bytes(b"image")
    sidecar_path = image_path.with_suffix(".txt")
    sidecar_path.write_text("keep me\n", encoding="utf-8")
    calls = _patch_jtp3_subprocess(tmp_path, monkeypatch)

    tag_jtp3.run_jtp3(tag_jtp3.Jtp3RunConfig(paths=[image_path]))

    assert capsys.readouterr().out == "from_side artist_(name)\n"
    assert sidecar_path.read_text(encoding="utf-8") == "keep me\n"
    command, _cwd, _env, _check, _text = calls["run"]
    assert isinstance(command, list)
    assert str(image_path) not in command


def test_run_jtp3_prints_paths_for_multiple_stdout_images(
    tmp_path: Path,
    monkeypatch,
    capsys,
) -> None:
    first_path = tmp_path / "a.png"
    second_path = tmp_path / "b.png"
    first_path.write_bytes(b"image")
    second_path.write_bytes(b"image")
    _patch_jtp3_subprocess(tmp_path, monkeypatch)

    tag_jtp3.run_jtp3(tag_jtp3.Jtp3RunConfig(paths=[first_path, second_path]))

    assert capsys.readouterr().out == (
        f"{first_path}\tfrom_side artist_(name)\n"
        f"{second_path}\tfrom_side artist_(name)\n"
    )


def test_run_jtp3_stdout_can_use_legacy_commas(
    tmp_path: Path,
    monkeypatch,
    capsys,
) -> None:
    image_path = tmp_path / "image.png"
    image_path.write_bytes(b"image")
    _patch_jtp3_subprocess(tmp_path, monkeypatch)

    tag_jtp3.run_jtp3(tag_jtp3.Jtp3RunConfig(paths=[image_path], comma_separated=True))

    assert capsys.readouterr().out == f"{UPSTREAM_TAGS}\n"


def test_run_jtp3_txt_postprocesses_original_sidecar(
    tmp_path: Path,
    monkeypatch,
) -> None:
    image_path = tmp_path / "image.png"
    image_path.write_bytes(b"image")
    _patch_jtp3_subprocess(tmp_path, monkeypatch)

    tag_jtp3.run_jtp3(tag_jtp3.Jtp3RunConfig(paths=[image_path], write_txt=True))

    assert image_path.with_suffix(".txt").read_text(encoding="utf-8") == (
        "from_side artist_(name)\n"
    )


def test_run_jtp3_txt_comma_preserves_legacy_sidecar(
    tmp_path: Path,
    monkeypatch,
) -> None:
    image_path = tmp_path / "image.png"
    image_path.write_bytes(b"image")
    _patch_jtp3_subprocess(tmp_path, monkeypatch)

    tag_jtp3.run_jtp3(
        tag_jtp3.Jtp3RunConfig(
            paths=[image_path],
            write_txt=True,
            comma_separated=True,
        )
    )

    assert image_path.with_suffix(".txt").read_text(encoding="utf-8") == (
        f"{UPSTREAM_TAGS}\n"
    )


def test_run_jtp3_rejects_csv_stdout_text_modes() -> None:
    with pytest.raises(ValueError, match="--csv-stdout"):
        tag_jtp3.run_jtp3(
            tag_jtp3.Jtp3RunConfig(
                paths=[Path("image.png")],
                csv_stdout=True,
                write_txt=True,
            )
        )


def test_parse_args_accepts_txt_and_comma_modes() -> None:
    args = tag_jtp3.parse_args(["--txt", "--comma", "image.png"])
    short_args = tag_jtp3.parse_args(["-c", "image.png"])

    assert args.txt is True
    assert args.comma_separated is True
    assert short_args.comma_separated is True


def test_ensure_jtp3_runtime_dependencies_reports_missing_pyvips(
    tmp_path: Path,
    monkeypatch,
) -> None:
    snapshot = tmp_path / "snapshot"
    snapshot.mkdir()
    (snapshot / "image.py").write_text("import pyvips\n", encoding="utf-8")

    def fake_import_module(name: str, package: str | None = None) -> object:
        raise ModuleNotFoundError("No module named 'pyvips'", name="pyvips")

    monkeypatch.setattr(tag_jtp3.importlib, "import_module", fake_import_module)

    with pytest.raises(tag_jtp3.Jtp3DependencyError, match="pyvips"):
        tag_jtp3.ensure_jtp3_runtime_dependencies(snapshot)


def test_ensure_jtp3_runtime_dependencies_reports_missing_libvips(
    tmp_path: Path,
    monkeypatch,
) -> None:
    snapshot = tmp_path / "snapshot"
    snapshot.mkdir()
    (snapshot / "image.py").write_text("import pyvips\n", encoding="utf-8")

    def fake_import_module(name: str, package: str | None = None) -> object:
        raise OSError("cannot load library 'vips-42'")

    monkeypatch.setattr(tag_jtp3.importlib, "import_module", fake_import_module)

    with pytest.raises(tag_jtp3.Jtp3DependencyError, match="libvips"):
        tag_jtp3.ensure_jtp3_runtime_dependencies(snapshot)


def test_ensure_jtp3_runtime_dependencies_skips_old_pillow_snapshot(
    tmp_path: Path,
    monkeypatch,
) -> None:
    snapshot = tmp_path / "snapshot"
    snapshot.mkdir()
    (snapshot / "image.py").write_text("from PIL import Image\n", encoding="utf-8")
    import_attempts: list[str] = []

    def fake_import_module(name: str, package: str | None = None) -> object:
        import_attempts.append(name)
        raise AssertionError("old snapshots should not import pyvips")

    monkeypatch.setattr(tag_jtp3.importlib, "import_module", fake_import_module)

    tag_jtp3.ensure_jtp3_runtime_dependencies(snapshot)

    assert import_attempts == []


def test_run_jtp3_downloads_limited_snapshot_and_runs_in_snapshot(
    tmp_path: Path,
    monkeypatch,
) -> None:
    snapshot = tmp_path / "snapshot"
    calls: dict[str, object] = {}

    def fake_snapshot_download(**kwargs: object) -> str:
        calls["snapshot_kwargs"] = kwargs
        return str(snapshot)

    def fake_run(
        command: list[str],
        *,
        cwd: Path,
        env: dict[str, str],
        check: bool,
        text: bool,
    ) -> subprocess.CompletedProcess[str]:
        calls["run"] = (command, cwd, env, check, text)
        return subprocess.CompletedProcess(command, 0)

    monkeypatch.setattr(tag_jtp3, "snapshot_download", fake_snapshot_download)
    monkeypatch.setattr(tag_jtp3.subprocess, "run", fake_run)

    result = tag_jtp3.run_jtp3(
        tag_jtp3.Jtp3RunConfig(paths=[Path("image.png")], csv_stdout=True),
        tag_jtp3.Jtp3SnapshotConfig(revision="abc123"),
    )

    assert result.returncode == 0
    snapshot_kwargs = calls["snapshot_kwargs"]
    assert isinstance(snapshot_kwargs, dict)
    assert snapshot_kwargs["repo_id"] == tag_jtp3.DEFAULT_REPO_ID
    assert snapshot_kwargs["revision"] == "abc123"
    assert snapshot_kwargs["allow_patterns"] == tag_jtp3.DEFAULT_ALLOW_PATTERNS
    command, cwd, env, check, text = calls["run"]
    assert isinstance(command, list)
    assert "--output" in command
    assert "-" in command
    assert cwd == snapshot
    assert isinstance(env, dict)
    assert env["PYTHONPATH"].split(os.pathsep)[0] == str(snapshot)
    assert check is True
    assert text is True