from __future__ import annotations

import json
from pathlib import Path

import pytest

from workflows import WorkflowBuilder, WorkflowError, fit_latent_size

BASE = Path(__file__).resolve().parents[1] / "workflows"


def make_builder() -> WorkflowBuilder:
    return WorkflowBuilder(
        BASE / "txt2img.json",
        BASE / "img2img.json",
        BASE / "mapping.json",
    )


def test_build_txt2img() -> None:
    builder = make_builder()
    workflow = builder.build_txt2img(prompt="cat", negative="bad", seed=123)

    assert workflow["67"]["inputs"]["text"] == "cat"
    assert workflow["71"]["inputs"]["text"] == "bad"
    assert workflow["70"]["inputs"]["seed"] == 123
    assert workflow["81"]["class_type"] == "EmptySD3LatentImage"
    assert workflow["70"]["inputs"]["denoise"] == 1.0


def test_build_img2img() -> None:
    builder = make_builder()
    workflow = builder.build_img2img(
        prompt="make blue",
        negative="bad",
        seed=7,
        image_name="input.png",
    )

    assert workflow["25"]["inputs"]["prompt"] == "make blue"
    assert workflow["25"]["inputs"]["vae"] == ["19", 2]
    assert workflow["10"]["inputs"]["text"] == "bad"
    assert workflow["3"]["inputs"]["noise_seed"] == 7
    assert workflow["3"]["inputs"]["steps"] == 12
    assert workflow["3"]["inputs"]["sampler_name"] == "dpmpp_2m"
    assert workflow["3"]["inputs"]["scheduler"] == "karras"
    assert workflow["20"]["class_type"] == "LoadImage"
    assert workflow["20"]["inputs"]["image"] == "input.png"
    assert workflow["15"]["inputs"]["width"] == 1024
    assert workflow["15"]["inputs"]["height"] == 1024


def test_build_img2img_custom_size() -> None:
    builder = make_builder()
    workflow = builder.build_img2img(
        prompt="make blue",
        negative="bad",
        seed=7,
        image_name="input.png",
        width=832,
        height=1216,
        steps=8,
        unet_name="Krea2-Q6.gguf",
    )

    assert workflow["15"]["inputs"]["width"] == 832
    assert workflow["15"]["inputs"]["height"] == 1216
    assert workflow["3"]["inputs"]["steps"] == 8
    assert workflow["19"]["inputs"]["unet_name"] == "Krea2-Q6.gguf"


def test_fit_latent_size_preserves_aspect_ratio() -> None:
    assert fit_latent_size(1024, 1024) == (1024, 1024)

    width, height = fit_latent_size(1080, 1920)
    assert width % 16 == 0
    assert height % 16 == 0
    assert height > width
    assert max(width, height) <= 1536

    width, height = fit_latent_size(4000, 3000)
    assert width % 16 == 0
    assert height % 16 == 0
    assert width > height


def test_fit_latent_size_fallback_for_invalid_size() -> None:
    assert fit_latent_size(0, 100) == (1024, 1024)


def test_templates_are_not_mutated() -> None:
    builder = make_builder()
    builder.build_txt2img(prompt="first", negative="bad", seed=1)
    second = builder.build_txt2img(prompt="second", negative="bad", seed=2)
    assert second["67"]["inputs"]["text"] == "second"


def test_invalid_mapping(tmp_path: Path) -> None:
    mapping = json.loads((BASE / "mapping.json").read_text(encoding="utf-8"))
    mapping["txt2img"]["positive"] = "999"
    path = tmp_path / "mapping.json"
    path.write_text(json.dumps(mapping), encoding="utf-8")

    with pytest.raises(WorkflowError):
        WorkflowBuilder(BASE / "txt2img.json", BASE / "img2img.json", path)