from __future__ import annotations

from pathlib import Path

from nbdraw.downloader import build_output_filename, download_image, infer_extension


class FakeResponse:
    def __init__(self, status_code=200, chunks=None):
        self.status_code = status_code
        self._chunks = chunks or [b"hello", b"world"]
        self.text = ""

    def iter_content(self, chunk_size=8192):
        del chunk_size
        yield from self._chunks


class FakeSession:
    def __init__(self):
        self.called_url = None

    def get(self, url, stream=True, timeout=None):
        del stream, timeout
        self.called_url = url
        return FakeResponse()


def test_infer_extension_fallback_png():
    assert infer_extension("https://a.com/file") == ".png"
    assert infer_extension("https://a.com/file.webp") == ".webp"


def test_build_output_filename_contains_task_id():
    name = build_output_filename("task-1")
    assert "task-1" in name
    assert name.endswith(".png")


def test_download_image_writes_file(tmp_path):
    session = FakeSession()
    path = download_image(
        "https://example.com/img.png",
        output_dir=str(tmp_path),
        task_id="task-xyz",
        session=session,
    )

    out = Path(path)
    assert out.exists()
    assert out.read_bytes() == b"helloworld"
    assert session.called_url == "https://example.com/img.png"
