import base64
from pathlib import Path

import pytest

from polza import PolzaClient, PolzaError


def test_reference_payload_accepts_url_and_data_uri():
    assert PolzaClient.reference_payload("https://example.com/style.png") == {
        "type": "url",
        "data": "https://example.com/style.png",
    }
    data_uri = "data:image/png;base64,AAAA"
    assert PolzaClient.reference_payload(data_uri) == {"type": "base64", "data": data_uri}


def test_reference_payload_encodes_local_image(tmp_path):
    path = tmp_path / "style.png"
    path.write_bytes(b"image-bytes")

    payload = PolzaClient.reference_payload(str(path))

    assert payload["type"] == "base64"
    assert payload["data"].startswith("data:image/png;base64,")
    assert base64.b64decode(payload["data"].split(",", 1)[1]) == b"image-bytes"


def test_reference_payload_rejects_unknown_path():
    with pytest.raises(PolzaError, match="does not exist"):
        PolzaClient.reference_payload("missing-style.png")


def test_create_image_sends_references_and_variants(monkeypatch):
    client = PolzaClient(api_key="test-key")
    captured = {}

    def fake_request(method, path, payload=None):
        captured.update(method=method, path=path, payload=payload)
        return {"id": "gen_1", "status": "pending"}

    monkeypatch.setattr(client, "_request", fake_request)
    client.create_image(
        model="seedream-3",
        prompt="game sprite",
        reference_images=["https://example.com/style.png"],
        count=3,
        aspect_ratio="1:1",
        seed=42,
        wait=False,
    )

    assert captured["method"] == "POST"
    assert captured["path"] == "/media"
    assert captured["payload"]["async"] is True
    assert captured["payload"]["input"]["max_images"] == 3
    assert captured["payload"]["input"]["images"][0]["type"] == "url"


def test_image_sources_handles_urls_and_base64():
    response = {
        "data": [{"url": "https://cdn.example/one.png"}],
        "result": {"images": [{"b64_json": "aGVsbG8="}]},
    }
    assert PolzaClient.image_sources(response) == [
        "https://cdn.example/one.png",
        "aGVsbG8=",
    ]