import httpx
import pytest
import respx

from anyfast import AnyFast


@pytest.fixture
def client():
    return AnyFast(api_key="test-key")


class TestImageGenerate:
    @respx.mock
    def test_generate(self, client):
        mock_response = {
            "created": 1700000000,
            "data": [{"url": "https://cdn.example.com/image.png"}],
        }

        route = respx.post("https://www.anyfast.ai/v1/images/generations").mock(
            return_value=httpx.Response(200, json=mock_response)
        )

        result = client.image.generate(
            model="flux-2-pro",
            prompt="a sunset",
            size="1024x1024",
        )

        assert route.called
        assert result["data"][0]["url"] == "https://cdn.example.com/image.png"

    @respx.mock
    def test_generate_seedream(self, client):
        respx.post("https://www.anyfast.ai/v1/images/generations").mock(
            return_value=httpx.Response(200, json={
                "id": "gen-001",
                "model": "doubao-seedream-5-0-260128",
                "status": "processing",
            })
        )

        result = client.image.generate(
            model="doubao-seedream-5-0-260128",
            prompt="a cat",
        )

        assert result["model"] == "doubao-seedream-5-0-260128"


class TestImageAsync:
    @respx.mock
    def test_generate_async(self, client):
        respx.post("https://www.anyfast.ai/v1/images/generations").mock(
            return_value=httpx.Response(200, json={"task_id": "img-task-001"})
        )

        response = client.image.generate_async(
            model="doubao-seedream-5-0-260128",
            prompt="a sunset",
        )

        assert response.task_id == "img-task-001"

    @respx.mock
    def test_query_task(self, client):
        respx.get("https://www.anyfast.ai/img-task-001").mock(
            return_value=httpx.Response(200, json={
                "id": "img-task-001",
                "status": "SUCCESS",
                "progress": "100%",
                "result": {"url": "https://cdn.example.com/result.png"},
            })
        )

        status = client.image.query_task("img-task-001")

        assert status.is_completed()
        assert status.result["url"] == "https://cdn.example.com/result.png"


class TestImageRun:
    @respx.mock
    def test_run(self, client):
        respx.post("https://www.anyfast.ai/v1/images/generations").mock(
            return_value=httpx.Response(200, json={"task_id": "img-run-001"})
        )
        respx.get("https://www.anyfast.ai/img-run-001").mock(
            side_effect=[
                httpx.Response(200, json={"id": "img-run-001", "status": "PROCESSING"}),
                httpx.Response(200, json={
                    "id": "img-run-001",
                    "status": "SUCCESS",
                    "result": {"url": "https://cdn.example.com/done.png"},
                }),
            ]
        )

        status = client.image.run(
            {"model": "flux-2-pro", "prompt": "a cat"},
            poll_interval=0.01,
        )

        assert status.is_completed()
        assert status.result["url"] == "https://cdn.example.com/done.png"

    @respx.mock
    def test_run_no_task_id(self, client):
        respx.post("https://www.anyfast.ai/v1/images/generations").mock(
            return_value=httpx.Response(200, json={"status": "ok"})
        )

        with pytest.raises(RuntimeError, match="No task ID"):
            client.image.run({"model": "test", "prompt": "test"})
