"""回収した tar の検査・展開のテスト（SPEC_03「pull」の二重の防御の 2 段目）.

受け取った tar にホワイトリスト外の名前・パス外に出る名前・リンクが 1 つでもあれば，
何も展開せずに中止すること（展開先に何も作られないこと）．
"""

import io
import tarfile
from pathlib import Path

import pytest

from tias_judge_model.remote.files import PushEntry
from tias_judge_model.remote.transfer import extract_pull_archive, write_tar_stream

RUN_ID = "20260927-120000_newton_abc1234"


def _add_file(tar: tarfile.TarFile, name: str, data: bytes = b"x") -> None:
    info = tarfile.TarInfo(name)
    info.size = len(data)
    tar.addfile(info, io.BytesIO(data))


def _add_link(tar: tarfile.TarFile, name: str, target: str, kind: bytes) -> None:
    info = tarfile.TarInfo(name)
    info.type = kind
    info.linkname = target
    tar.addfile(info)


def _archive(build) -> io.BytesIO:
    buffer = io.BytesIO()
    with tarfile.open(fileobj=buffer, mode="w", format=tarfile.PAX_FORMAT) as tar:
        build(tar)
    buffer.seek(0)
    return buffer


def _all_files(root: Path) -> list[Path]:
    return sorted(root.rglob("*")) if root.exists() else []


def test_extracts_whitelisted_files(tmp_path: Path) -> None:
    def build(tar: tarfile.TarFile) -> None:
        _add_file(tar, f"{RUN_ID}/config.yaml", b"seed: 0\n")
        _add_file(tar, f"{RUN_ID}/metrics.json", b"{}")
        _add_file(tar, f"{RUN_ID}/figures/confusion_matrix.png", b"png")

    dest = tmp_path / "runs"
    result = extract_pull_archive(_archive(build), dest, RUN_ID, with_checkpoints=False)
    assert result.files == 3
    assert (dest / RUN_ID / "config.yaml").read_bytes() == b"seed: 0\n"
    assert (dest / RUN_ID / "figures" / "confusion_matrix.png").read_bytes() == b"png"


def test_overwrites_existing_file(tmp_path: Path) -> None:
    dest = tmp_path / "runs"
    (dest / RUN_ID).mkdir(parents=True)
    (dest / RUN_ID / "metrics.json").write_bytes(b"old")
    archive = _archive(lambda tar: _add_file(tar, f"{RUN_ID}/metrics.json", b"new"))
    extract_pull_archive(archive, dest, RUN_ID, with_checkpoints=False)
    assert (dest / RUN_ID / "metrics.json").read_bytes() == b"new"


def test_checkpoints_extracted_only_with_flag(tmp_path: Path) -> None:
    def build(tar: tarfile.TarFile) -> None:
        _add_file(tar, f"{RUN_ID}/config.yaml")
        _add_file(tar, f"{RUN_ID}/checkpoints/best.pt", b"weights")

    dest = tmp_path / "runs"
    with pytest.raises(ValueError):
        extract_pull_archive(_archive(build), dest, RUN_ID, with_checkpoints=False)
    assert _all_files(dest) == []

    extract_pull_archive(_archive(build), dest, RUN_ID, with_checkpoints=True)
    assert (dest / RUN_ID / "checkpoints" / "best.pt").read_bytes() == b"weights"


BAD_MEMBERS = {
    "split_json": lambda tar: _add_file(tar, f"{RUN_ID}/split.json"),
    "predictions": lambda tar: _add_file(tar, f"{RUN_ID}/predictions.csv"),
    "png_outside_figures": lambda tar: _add_file(tar, f"{RUN_ID}/P0001.png"),
    "manifest": lambda tar: _add_file(tar, f"{RUN_ID}/manifest.csv"),
    "image": lambda tar: _add_file(tar, f"{RUN_ID}/images/P0001.jpg"),
    "other_run": lambda tar: _add_file(tar, "other_run/config.yaml"),
    "dotdot": lambda tar: _add_file(tar, f"{RUN_ID}/../evil.txt"),
    "dotdot_top": lambda tar: _add_file(tar, "../evil.txt"),
    "absolute": lambda tar: _add_file(tar, "/tmp/evil.txt"),
    "drive_letter": lambda tar: _add_file(tar, "C:/evil.txt"),
    "drive_relative": lambda tar: _add_file(tar, "C:evil.txt"),
    "backslash": lambda tar: _add_file(tar, f"{RUN_ID}\\config.yaml"),
    "symlink": lambda tar: _add_link(tar, f"{RUN_ID}/log.txt", "/etc/passwd", tarfile.SYMTYPE),
    "symlink_relative": lambda tar: _add_link(
        tar, f"{RUN_ID}/history.csv", "../../outside", tarfile.SYMTYPE
    ),
    "hardlink": lambda tar: _add_link(
        tar, f"{RUN_ID}/metrics.json", f"{RUN_ID}/config.yaml", tarfile.LNKTYPE
    ),
}


@pytest.mark.parametrize("bad", list(BAD_MEMBERS), ids=list(BAD_MEMBERS))
@pytest.mark.parametrize("position", ["first", "last"])
def test_bad_member_aborts_without_extracting_anything(
    tmp_path: Path, bad: str, position: str
) -> None:
    def build(tar: tarfile.TarFile) -> None:
        if position == "first":
            BAD_MEMBERS[bad](tar)
        _add_file(tar, f"{RUN_ID}/config.yaml")
        _add_file(tar, f"{RUN_ID}/figures/cm.png")
        if position == "last":
            BAD_MEMBERS[bad](tar)

    work = tmp_path / "work"
    dest = work / "runs"
    with pytest.raises(ValueError):
        extract_pull_archive(_archive(build), dest, RUN_ID, with_checkpoints=True)
    # 展開先にも，その外側（パス外に出る名前の行き先）にも何も作られない
    assert _all_files(dest) == []
    assert not (dest / RUN_ID / "config.yaml").exists()
    assert not (work / "evil.txt").exists()
    assert not (tmp_path / "evil.txt").exists()
    assert [p for p in _all_files(tmp_path) if p.is_file()] == []


def test_error_message_does_not_contain_rejected_names(tmp_path: Path) -> None:
    secret = "P0001_yamada"

    def build(tar: tarfile.TarFile) -> None:
        _add_file(tar, f"{RUN_ID}/config.yaml")
        _add_file(tar, f"{RUN_ID}/{secret}.png")

    with pytest.raises(ValueError) as excinfo:
        extract_pull_archive(_archive(build), tmp_path / "runs", RUN_ID, with_checkpoints=False)
    assert secret not in str(excinfo.value)


def test_write_tar_stream_round_trip(tmp_path: Path) -> None:
    src = tmp_path / "src"
    (src / "pkg").mkdir(parents=True)
    (src / "pkg" / "a.py").write_bytes(b"print('a')\n")
    (src / "日本語.md").write_text("x", encoding="utf-8")
    entries = [
        PushEntry(src / "pkg" / "a.py", "pkg/a.py", (src / "pkg" / "a.py").stat().st_size),
        PushEntry(src / "日本語.md", "日本語.md", 1),
    ]
    buffer = io.BytesIO()
    write_tar_stream(entries, buffer)
    buffer.seek(0)
    with tarfile.open(fileobj=buffer, mode="r") as tar:
        names = tar.getnames()
        assert names == ["pkg/a.py", "日本語.md"]
        member = tar.extractfile("pkg/a.py")
        assert member is not None
        assert member.read() == b"print('a')\n"