from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Callable, Sequence
from pathlib import Path

import pytest

import agents.sandbox.session.manifest_application as manifest_application_module
from agents.sandbox.entries import (
    Dir,
    File,
    GCSMount,
    InContainerMountStrategy,
    MountpointMountPattern,
)
from agents.sandbox.errors import ExecNonZeroError
from agents.sandbox.manifest import Manifest
from agents.sandbox.materialization import MaterializedFile
from agents.sandbox.session.manifest_application import ManifestApplier
from agents.sandbox.types import ExecResult, Group, User


def _materialized(dest: Path) -> list[MaterializedFile]:
    return [MaterializedFile(path=dest, sha256=dest.as_posix())]


@pytest.mark.asyncio
async def test_manifest_applier_only_applies_ephemeral_entries_without_account_provisioning() -> (
    None
):
    mkdir_calls: list[Path] = []
    exec_calls: list[tuple[str, ...]] = []
    apply_calls: list[tuple[str, Path, Path]] = []

    async def mkdir(path: Path) -> None:
        mkdir_calls.append(path)

    async def exec_checked_nonzero(*command: str) -> ExecResult:
        exec_calls.append(command)
        return ExecResult(stdout=b"", stderr=b"", exit_code=0)

    async def apply_entry(entry: object, dest: Path, base_dir: Path) -> list[MaterializedFile]:
        apply_calls.append((type(entry).__name__, dest, base_dir))
        return _materialized(dest)

    applier = ManifestApplier(
        mkdir=mkdir,
        exec_checked_nonzero=exec_checked_nonzero,
        apply_entry=apply_entry,
    )
    manifest = Manifest(
        root="/workspace",
        entries={
            "keep.txt": File(content=b"keep"),
            "tmp.txt": File(content=b"tmp", ephemeral=True),
        },
        users=[User(name="alice")],
        groups=[Group(name="dev", users=[User(name="alice")])],
    )

    result = await applier.apply_manifest(manifest, only_ephemeral=True)

    assert mkdir_calls == [Path("/workspace")]
    assert exec_calls == []
    assert apply_calls == [("File", Path("/workspace/tmp.txt"), Path("/"))]
    assert result.files == _materialized(Path("/workspace/tmp.txt"))


@pytest.mark.asyncio
async def test_manifest_applier_only_ephemeral_reapplies_nested_ephemeral_children() -> None:
    apply_calls: list[tuple[str, Path, Path]] = []

    async def mkdir(_path: Path) -> None:
        return None

    async def exec_checked_nonzero(*_command: str) -> ExecResult:
        return ExecResult(stdout=b"", stderr=b"", exit_code=0)

    async def apply_entry(entry: object, dest: Path, base_dir: Path) -> list[MaterializedFile]:
        apply_calls.append((type(entry).__name__, dest, base_dir))
        return _materialized(dest)

    applier = ManifestApplier(
        mkdir=mkdir,
        exec_checked_nonzero=exec_checked_nonzero,
        apply_entry=apply_entry,
    )
    manifest = Manifest(
        root="/workspace",
        entries={
            "dir": Dir(
                children={
                    "keep.txt": File(content=b"keep"),
                    "tmp.txt": File(content=b"tmp", ephemeral=True),
                }
            )
        },
    )

    result = await applier.apply_manifest(manifest, only_ephemeral=True)

    assert apply_calls == [("File", Path("/workspace/dir/tmp.txt"), Path("/"))]
    assert result.files == _materialized(Path("/workspace/dir/tmp.txt"))


@pytest.mark.asyncio
async def test_manifest_applier_only_ephemeral_reapplies_full_ephemeral_directories() -> None:
    applied_entries: list[tuple[object, Path, Path]] = []

    async def mkdir(_path: Path) -> None:
        return None

    async def exec_checked_nonzero(*_command: str) -> ExecResult:
        return ExecResult(stdout=b"", stderr=b"", exit_code=0)

    async def apply_entry(entry: object, dest: Path, base_dir: Path) -> list[MaterializedFile]:
        applied_entries.append((entry, dest, base_dir))
        return _materialized(dest)

    applier = ManifestApplier(
        mkdir=mkdir,
        exec_checked_nonzero=exec_checked_nonzero,
        apply_entry=apply_entry,
    )
    manifest = Manifest(
        root="/workspace",
        entries={
            "tmp": Dir(
                ephemeral=True,
                children={
                    "keep.txt": File(content=b"keep"),
                    "nested": Dir(children={"child.txt": File(content=b"child")}),
                    "tmp.txt": File(content=b"tmp", ephemeral=True),
                },
            )
        },
    )

    result = await applier.apply_manifest(manifest, only_ephemeral=True)

    assert len(applied_entries) == 1
    entry, dest, base_dir = applied_entries[0]
    assert isinstance(entry, Dir)
    assert dest == Path("/workspace/tmp")
    assert base_dir == Path("/")
    assert set(entry.children) == {"keep.txt", "nested", "tmp.txt"}
    assert result.files == _materialized(Path("/workspace/tmp"))


@pytest.mark.asyncio
async def test_manifest_applier_respects_explicit_base_dir() -> None:
    apply_calls: list[tuple[str, Path, Path]] = []

    async def mkdir(_path: Path) -> None:
        return None

    async def exec_checked_nonzero(*_command: str) -> ExecResult:
        return ExecResult(stdout=b"", stderr=b"", exit_code=0)

    async def apply_entry(entry: object, dest: Path, base_dir: Path) -> list[MaterializedFile]:
        apply_calls.append((type(entry).__name__, dest, base_dir))
        return _materialized(dest)

    applier = ManifestApplier(
        mkdir=mkdir,
        exec_checked_nonzero=exec_checked_nonzero,
        apply_entry=apply_entry,
    )
    manifest = Manifest(entries={"file.txt": File(content=b"hello")})

    result = await applier.apply_manifest(manifest, base_dir=Path("/tmp/project"))

    assert apply_calls == [("File", Path("/workspace/file.txt"), Path("/tmp/project"))]
    assert result.files == _materialized(Path("/workspace/file.txt"))


@pytest.mark.asyncio
async def test_manifest_applier_caps_parallel_entry_batch(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    observed_limits: list[int | None] = []

    async def gather_with_limit_recording(
        task_factories: Sequence[Callable[[], Awaitable[list[MaterializedFile]]]],
        *,
        max_concurrency: int | None = None,
    ) -> list[list[MaterializedFile]]:
        observed_limits.append(max_concurrency)
        return [await factory() for factory in task_factories]

    async def mkdir(_path: Path) -> None:
        return None

    async def exec_checked_nonzero(*_command: str) -> ExecResult:
        return ExecResult(stdout=b"", stderr=b"", exit_code=0)

    async def apply_entry(_entry: object, dest: Path, _base_dir: Path) -> list[MaterializedFile]:
        return _materialized(dest)

    monkeypatch.setattr(
        manifest_application_module,
        "gather_in_order",
        gather_with_limit_recording,
    )
    applier = ManifestApplier(
        mkdir=mkdir,
        exec_checked_nonzero=exec_checked_nonzero,
        apply_entry=apply_entry,
        max_entry_concurrency=2,
    )

    result = await applier.apply_manifest(
        Manifest(entries={"a.txt": File(content=b"a"), "b.txt": File(content=b"b")})
    )

    assert observed_limits == [2]
    assert result.files == [
        MaterializedFile(path=Path("/workspace/a.txt"), sha256="/workspace/a.txt"),
        MaterializedFile(path=Path("/workspace/b.txt"), sha256="/workspace/b.txt"),
    ]


@pytest.mark.asyncio
async def test_manifest_applier_provisions_groups_and_unique_users_before_entries() -> None:
    exec_calls: list[tuple[str, ...]] = []

    async def mkdir(_path: Path) -> None:
        return None

    async def exec_checked_nonzero(*command: str) -> ExecResult:
        exec_calls.append(command)
        return ExecResult(stdout=b"", stderr=b"", exit_code=0)

    async def apply_entry(_entry: object, _dest: Path, _base_dir: Path) -> list[MaterializedFile]:
        return []

    applier = ManifestApplier(
        mkdir=mkdir,
        exec_checked_nonzero=exec_checked_nonzero,
        apply_entry=apply_entry,
    )
    manifest = Manifest(
        users=[User(name="alice")],
        groups=[Group(name="dev", users=[User(name="alice"), User(name="bob")])],
    )

    result = await applier.apply_manifest(manifest)

    assert result.files == []
    assert exec_calls[0] == ("groupadd", "dev")
    assert exec_calls.count(("groupadd", "alice")) == 0
    assert exec_calls.count(("groupadd", "bob")) == 0
    assert ("useradd", "-U", "-M", "-s", "/usr/sbin/nologin", "alice") in exec_calls
    assert ("useradd", "-U", "-M", "-s", "/usr/sbin/nologin", "bob") in exec_calls
    assert ("usermod", "-aG", "dev", "alice") in exec_calls
    assert ("usermod", "-aG", "dev", "bob") in exec_calls


@pytest.mark.asyncio
async def test_manifest_applier_can_apply_full_manifest_without_account_provisioning() -> None:
    exec_calls: list[tuple[str, ...]] = []
    apply_calls: list[tuple[str, Path, Path]] = []

    async def mkdir(_path: Path) -> None:
        return None

    async def exec_checked_nonzero(*command: str) -> ExecResult:
        exec_calls.append(command)
        return ExecResult(stdout=b"", stderr=b"", exit_code=0)

    async def apply_entry(entry: object, dest: Path, base_dir: Path) -> list[MaterializedFile]:
        apply_calls.append((type(entry).__name__, dest, base_dir))
        return _materialized(dest)

    applier = ManifestApplier(
        mkdir=mkdir,
        exec_checked_nonzero=exec_checked_nonzero,
        apply_entry=apply_entry,
    )
    manifest = Manifest(
        entries={"file.txt": File(content=b"hello")},
        users=[User(name="alice")],
        groups=[Group(name="dev", users=[User(name="alice")])],
    )

    result = await applier.apply_manifest(manifest, provision_accounts=False)

    assert exec_calls == []
    assert apply_calls == [("File", Path("/workspace/file.txt"), Path("/"))]
    assert result.files == _materialized(Path("/workspace/file.txt"))


@pytest.mark.asyncio
async def test_manifest_applier_raises_with_command_stdout_and_stderr_on_provision_failure() -> (
    None
):
    async def mkdir(_path: Path) -> None:
        return None

    async def exec_checked_nonzero(*command: str) -> ExecResult:
        raise ExecNonZeroError(
            ExecResult(stdout=b"groupadd output", stderr=b"groupadd failed", exit_code=9),
            command=command,
        )

    async def apply_entry(_entry: object, _dest: Path, _base_dir: Path) -> list[MaterializedFile]:
        return []

    applier = ManifestApplier(
        mkdir=mkdir,
        exec_checked_nonzero=exec_checked_nonzero,
        apply_entry=apply_entry,
    )
    manifest = Manifest(groups=[Group(name="dev", users=[])])

    with pytest.raises(ExecNonZeroError) as exc_info:
        await applier.apply_manifest(manifest)

    assert exc_info.value.context["command"] == ("groupadd", "dev")
    assert exc_info.value.context["command_str"] == "groupadd dev"
    assert exc_info.value.context["stdout"] == "groupadd output"
    assert exc_info.value.context["stderr"] == "groupadd failed"
    assert exc_info.value.message == "stdout: groupadd output\nstderr: groupadd failed"


@pytest.mark.asyncio
async def test_manifest_applier_raises_without_stream_labels_when_only_stdout_is_present() -> None:
    async def mkdir(_path: Path) -> None:
        return None

    async def exec_checked_nonzero(*command: str) -> ExecResult:
        raise ExecNonZeroError(
            ExecResult(stdout=b"useradd unavailable", stderr=b"", exit_code=127),
            command=command,
        )

    async def apply_entry(_entry: object, _dest: Path, _base_dir: Path) -> list[MaterializedFile]:
        return []

    applier = ManifestApplier(
        mkdir=mkdir,
        exec_checked_nonzero=exec_checked_nonzero,
        apply_entry=apply_entry,
    )
    manifest = Manifest(users=[User(name="sandbox-user")])

    with pytest.raises(ExecNonZeroError) as exc_info:
        await applier.apply_manifest(manifest)

    assert exc_info.value.context["command_str"] == (
        "useradd -U -M -s /usr/sbin/nologin sandbox-user"
    )
    assert exc_info.value.context["stdout"] == "useradd unavailable"
    assert exc_info.value.context["stderr"] == ""
    assert exc_info.value.message == "useradd unavailable"


@pytest.mark.asyncio
async def test_apply_entry_batch_flushes_parallel_work_before_overlapping_paths() -> None:
    events: list[tuple[str, Path]] = []

    async def mkdir(_path: Path) -> None:
        return None

    async def exec_checked_nonzero(*_command: str) -> ExecResult:
        return ExecResult(stdout=b"", stderr=b"", exit_code=0)

    async def apply_entry(_entry: object, dest: Path, _base_dir: Path) -> list[MaterializedFile]:
        events.append(("start", dest))
        await asyncio.sleep(0)
        events.append(("end", dest))
        return _materialized(dest)

    applier = ManifestApplier(
        mkdir=mkdir,
        exec_checked_nonzero=exec_checked_nonzero,
        apply_entry=apply_entry,
    )
    destinations = [
        Path("/workspace/alpha.txt"),
        Path("/workspace/beta.txt"),
        Path("/workspace/nested"),
        Path("/workspace/nested/child.txt"),
    ]

    files = await applier._apply_entry_batch(
        [
            (destinations[0], File(content=b"a")),
            (destinations[1], File(content=b"b")),
            (destinations[2], Dir()),
            (destinations[3], File(content=b"c")),
        ],
        base_dir=Path("/"),
    )

    assert [file.path for file in files] == destinations
    child_start = events.index(("start", destinations[3]))
    assert events.index(("end", destinations[0])) < child_start
    assert events.index(("end", destinations[1])) < child_start
    assert events.index(("end", destinations[2])) < child_start


@pytest.mark.asyncio
async def test_apply_entry_batch_flushes_before_and_after_mount_entries() -> None:
    events: list[tuple[str, Path]] = []

    async def mkdir(_path: Path) -> None:
        return None

    async def exec_checked_nonzero(*_command: str) -> ExecResult:
        return ExecResult(stdout=b"", stderr=b"", exit_code=0)

    async def apply_entry(_entry: object, dest: Path, _base_dir: Path) -> list[MaterializedFile]:
        events.append(("start", dest))
        await asyncio.sleep(0)
        events.append(("end", dest))
        return _materialized(dest)

    applier = ManifestApplier(
        mkdir=mkdir,
        exec_checked_nonzero=exec_checked_nonzero,
        apply_entry=apply_entry,
    )
    destinations = [
        Path("/workspace/alpha.txt"),
        Path("/workspace/beta.txt"),
        Path("/workspace/mount"),
        Path("/workspace/gamma.txt"),
    ]

    files = await applier._apply_entry_batch(
        [
            (destinations[0], File(content=b"a")),
            (destinations[1], File(content=b"b")),
            (
                destinations[2],
                GCSMount(
                    bucket="sandbox-bucket",
                    mount_strategy=InContainerMountStrategy(pattern=MountpointMountPattern()),
                ),
            ),
            (destinations[3], File(content=b"c")),
        ],
        base_dir=Path("/"),
    )

    assert [file.path for file in files] == destinations
    mount_start = events.index(("start", destinations[2]))
    gamma_start = events.index(("start", destinations[3]))
    assert events.index(("end", destinations[0])) < mount_start
    assert events.index(("end", destinations[1])) < mount_start
    assert events.index(("end", destinations[2])) < gamma_start