from io import BytesIO
from pathlib import Path

import pytest

from prodigy_teams.commands.files.cp import _download
from prodigy_teams_broker_sdk import models as broker_models


class _FakeFilesAPI:
    def __init__(self, expected_src: str) -> None:
        self.expected_src = expected_src
        self.downloaded_targets: list[str] = []

    def list_dir(self, body: broker_models.Listing) -> broker_models.PathList:
        assert body.path == self.expected_src
        return broker_models.PathList(
            paths=[
                "/mnt/nfs/data/agent-workflow-tests/model/model-best/config.cfg",
                "/mnt/nfs/data/agent-workflow-tests/model/model-best/meta.json",
            ],
            stats=None,
            exists=True,
            is_file=False,
        )

    def download(self, body: broker_models.Downloading) -> BytesIO:
        self.downloaded_targets.append(body.target)
        return BytesIO(body.target.encode("utf8"))


class _FakeBrokerClient:
    def __init__(self, expected_src: str) -> None:
        self.files = _FakeFilesAPI(expected_src)


@pytest.mark.parametrize(
    "src",
    [
        "{__data__}/agent-workflow-tests/model",
        "{__nfs__}/data/agent-workflow-tests/model",
    ],
)
def test_download_directory_from_builtin_alias_uses_requested_root(
    src: str,
    tmp_path: Path,
) -> None:
    broker_client = _FakeBrokerClient(src)
    dest = tmp_path / "downloads"
    dest.mkdir()

    _download(
        src=src,
        dest=dest,
        make_dirs=False,
        overwrite=False,
        recurse=True,
        broker_client=broker_client,  # type: ignore[arg-type]
    )

    assert (
        dest / "model" / "model-best" / "config.cfg"
    ).read_text() == "/mnt/nfs/data/agent-workflow-tests/model/model-best/config.cfg"
    assert (
        dest / "model" / "model-best" / "meta.json"
    ).read_text() == "/mnt/nfs/data/agent-workflow-tests/model/model-best/meta.json"
