import base64
import json

import httpx
import pytest
import respx

from tests.conftest import COMFY, png_bytes, png_size

GEN = "http://adapter.test/v1/images/generations"
PROMPT_ID = "p-1"


def _history_json(width=64, height=64):
    return {
        PROMPT_ID: {
            "outputs": {"60": {"images": [{"filename": "out_00001_.png", "subfolder": "", "type": "output"}]}},
            "status": {"status_str": "success", "completed": True},
        }
    }


def _mock_pipeline(view_png: bytes | None = None, prompt_response: httpx.Response | None = None):
    prompt_route = respx.post(f"{COMFY}/prompt").mock(
        return_value=prompt_response or httpx.Response(200, json={"prompt_id": PROMPT_ID})
    )
    respx.get(f"{COMFY}/history/{PROMPT_ID}").mock(return_value=httpx.Response(200, json=_history_json()))
    view_response = httpx.Response(200, content=view_png) if view_png else httpx.Response(500)
    respx.get(f"{COMFY}/view").mock(return_value=view_response)
    return prompt_route


@respx.mock
async def test_generation_success(client: httpx.AsyncClient):
    png = png_bytes(1024, 1024)
    prompt_route = _mock_pipeline(view_png=png)
    r = await client.post(GEN, json={"prompt": "a red cube", "model": "z-image-turbo", "size": "1024x1024", "seed": 123})
    assert r.status_code == 200
    body = r.json()
    assert set(body) == {"created", "data"} and set(body["data"][0]) == {"b64_json"}
    img = base64.b64decode(body["data"][0]["b64_json"])
    assert png_size(img) == (1024, 1024)
    # injected values reached the workflow
    sent = json.loads(prompt_route.calls.last.request.content.decode())["prompt"]
    assert sent["83:27"]["inputs"]["text"] == "a red cube"
    assert sent["83:13"]["inputs"]["width"] == 1024 and sent["83:13"]["inputs"]["height"] == 1024
    assert sent["83:3"]["inputs"]["seed"] == 123


@respx.mock
async def test_generation_size_custom(client: httpx.AsyncClient):
    prompt_route = _mock_pipeline(view_png=png_bytes(832, 1216))
    r = await client.post(GEN, json={"prompt": "x", "size": "832x1216"})
    assert r.status_code == 200
    sent = json.loads(prompt_route.calls.last.request.content.decode())["prompt"]
    assert (sent["83:13"]["inputs"]["width"], sent["83:13"]["inputs"]["height"]) == (832, 1216)


@respx.mock
async def test_generation_seed_random_per_request(client: httpx.AsyncClient):
    prompt_route = _mock_pipeline(view_png=png_bytes())
    await client.post(GEN, json={"prompt": "x"})
    seed_a = json.loads(prompt_route.calls.last.request.content.decode())["prompt"]["83:3"]["inputs"]["seed"]
    await client.post(GEN, json={"prompt": "x"})
    seed_b = json.loads(prompt_route.calls.last.request.content.decode())["prompt"]["83:3"]["inputs"]["seed"]
    assert seed_a != seed_b


async def _assert_400(r):
    assert r.status_code == 400
    err = r.json()["error"]
    assert err["type"] == "invalid_request_error" and err["message"]


async def test_empty_prompt_400(client):
    await _assert_400(await client.post(GEN, json={"prompt": "   "}))


async def test_n_gt_1_400(client):
    await _assert_400(await client.post(GEN, json={"prompt": "x", "n": 2}))


@pytest.mark.parametrize("bad", ["wide", "1024", "1024x", "0x512"])
async def test_bad_size_400(client, bad):
    await _assert_400(await client.post(GEN, json={"prompt": "x", "size": bad}))


async def test_unknown_model_400(client):
    await _assert_400(await client.post(GEN, json={"prompt": "x", "model": "dall-e-3"}))


async def test_bad_json_400(client):
    await _assert_400(await client.post(GEN, content=b"{not json", headers={"Content-Type": "application/json"}))


@respx.mock
async def test_comfy_validation_failure_502(client):
    _mock_pipeline(prompt_response=httpx.Response(400, json={"node_errors": {"83:27": {"errors": [{"message": "bad", "details": "'83:43'"}]}}}))
    r = await client.post(GEN, json={"prompt": "x"})
    assert r.status_code == 502
    assert r.json()["error"]["code"] == "comfyui_validation_failed"


@respx.mock
async def test_comfy_execution_failure_502(client):
    respx.post(f"{COMFY}/prompt").mock(return_value=httpx.Response(200, json={"prompt_id": PROMPT_ID}))
    respx.get(f"{COMFY}/history/{PROMPT_ID}").mock(
        return_value=httpx.Response(200, json={PROMPT_ID: {"outputs": {}, "status": {"status_str": "error", "completed": False}}})
    )
    r = await client.post(GEN, json={"prompt": "x"})
    assert r.status_code == 502
    assert r.json()["error"]["code"] == "comfyui_execution_failed"


@respx.mock
async def test_comfy_timeout_502(client):
    respx.post(f"{COMFY}/prompt").mock(return_value=httpx.Response(200, json={"prompt_id": PROMPT_ID}))
    respx.get(f"{COMFY}/history/{PROMPT_ID}").mock(return_value=httpx.Response(200, json={}))
    # client fixture uses execution_timeout=5; shorten via a second app instance instead:
    import contextlib

    from app.main import create_app
    from tests.conftest import settings as make_settings

    app = create_app(make_settings(execution_timeout=0.05, poll_interval=0.0))
    async with contextlib.AsyncExitStack() as stack:
        await stack.enter_async_context(app.router.lifespan_context(app))
        transport = httpx.ASGITransport(app=app)
        ac = await stack.enter_async_context(httpx.AsyncClient(transport=transport, base_url="http://adapter.test"))
        ac.headers["Authorization"] = client.headers["Authorization"]
        r = await ac.post(GEN, json={"prompt": "x"})
    assert r.status_code == 502
    assert r.json()["error"]["code"] == "comfyui_timeout"


@respx.mock
async def test_comfy_unreachable_502(client):
    respx.post(f"{COMFY}/prompt").mock(side_effect=httpx.ConnectError("boom"))
    r = await client.post(GEN, json={"prompt": "x"})
    assert r.status_code == 502