from __future__ import annotations

import importlib.util
from pathlib import Path

import torch

PLUGIN_ROOT = Path(__file__).resolve().parents[1]


def load_models_module():
    module_path = PLUGIN_ROOT / "lib" / "models.py"
    assert module_path.exists(), f"Expected module at {module_path}"
    spec = importlib.util.spec_from_file_location("comfyui_kieapi_models", module_path)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def load_image_utils_module():
    module_path = PLUGIN_ROOT / "lib" / "image_utils.py"
    assert module_path.exists(), f"Expected module at {module_path}"
    spec = importlib.util.spec_from_file_location("comfyui_kieapi_image_utils", module_path)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def test_upload_request_serializes_expected_fields():
    module = load_models_module()

    request = module.KieBase64UploadRequest(
        base64Data="data:image/png;base64,abc123",
        uploadPath="images/base64",
        fileName="reference.png",
    )

    payload = request.model_dump()
    assert payload["base64Data"] == "data:image/png;base64,abc123"
    assert payload["uploadPath"] == "images/base64"
    assert payload["fileName"] == "reference.png"


def test_create_task_request_serializes_expected_model_and_input():
    module = load_models_module()

    request = module.KieCreateTaskRequest(
        input=module.KieNanoBanana2Input(
            prompt="banana ship",
            aspect_ratio="auto",
            resolution="2K",
            output_format="jpg",
            google_search=False,
            image_input=["https://example.com/reference.png"],
        )
    )

    payload = request.model_dump()
    assert payload["model"] == "nano-banana-2"
    assert payload["input"]["prompt"] == "banana ship"
    assert payload["input"]["image_input"] == ["https://example.com/reference.png"]


def test_task_detail_parses_result_urls_from_result_json():
    module = load_models_module()

    response = module.KieTaskRecordResponse.model_validate(
        {
            "code": 200,
            "msg": "success",
            "data": {
                "taskId": "task_123",
                "state": "success",
                "resultJson": "{\"resultUrls\":[\"https://example.com/1.jpg\",\"https://example.com/2.jpg\"]}",
                "failCode": "",
                "failMsg": "",
            },
        }
    )

    assert response.data.result_urls() == ["https://example.com/1.jpg", "https://example.com/2.jpg"]


def test_task_detail_exposes_failure_information():
    module = load_models_module()

    response = module.KieTaskRecordResponse.model_validate(
        {
            "code": 200,
            "msg": "success",
            "data": {
                "taskId": "task_123",
                "state": "fail",
                "resultJson": "",
                "failCode": "content_blocked",
                "failMsg": "prompt rejected",
            },
        }
    )

    assert response.data.failure_reason() == "content_blocked: prompt rejected"


def test_tensor_to_base64_png_returns_data_uri():
    module = load_image_utils_module()
    image = torch.tensor([[[[1.0, 0.0, 0.0]]]], dtype=torch.float32)

    encoded = module.tensor_to_base64_png(image)

    assert encoded.startswith("data:image/png;base64,")


def test_tensor_batch_to_base64_pngs_returns_one_data_uri_per_image():
    module = load_image_utils_module()
    image_batch = torch.tensor(
        [
            [[[1.0, 0.0, 0.0]]],
            [[[0.0, 1.0, 0.0]]],
            [[[0.0, 0.0, 1.0]]],
        ],
        dtype=torch.float32,
    )

    encoded = module.tensor_batch_to_base64_pngs(image_batch)

    assert len(encoded) == 3
    assert all(item.startswith("data:image/png;base64,") for item in encoded)


def test_decode_image_bytes_returns_single_image_tensor():
    module = load_image_utils_module()
    image = torch.tensor([[[[0.0, 1.0, 0.0]]]], dtype=torch.float32)
    encoded = module.tensor_to_base64_png(image)
    raw_bytes = module.base64_data_uri_to_bytes(encoded)

    decoded = module.decode_image_bytes(raw_bytes)

    assert tuple(decoded.shape) == (1, 1, 1, 3)
    assert torch.isclose(decoded[0, 0, 0, 1], torch.tensor(1.0))


def test_stack_image_tensors_batches_multiple_results():
    module = load_image_utils_module()
    first = torch.zeros((1, 2, 2, 3), dtype=torch.float32)
    second = torch.ones((1, 2, 2, 3), dtype=torch.float32)

    batch = module.stack_image_tensors([first, second])

    assert tuple(batch.shape) == (2, 2, 2, 3)
    assert torch.equal(batch[0], first[0])
    assert torch.equal(batch[1], second[0])