import base64

import cv2
import numpy as np
import pytest
from fastapi.testclient import TestClient

import app as worker


@pytest.mark.parametrize(("densities", "labels"), [
    (["dense"], ["适中偏多"]),
    (["sparse", "normal", "dense"], ["稀疏", "正常", "适中偏多"]),
    ([], []),
])
def test_interface22_selects_prompts_and_aligns_output_size(monkeypatch, densities, labels):
    from hairline import service

    prompts = []
    rounded_png = cv2.imencode(".png", np.zeros((15, 13, 3), np.uint8))[1].tobytes()

    def generate(image, styles, defaults, *, prompt, hairline_height, **kwargs):
        assert styles == [1, 3]
        assert hairline_height == -1
        prompts.append(prompt)
        return [
            {"order": order, "hairline_type": key, "grown_png": rounded_png,
             "overlay": np.zeros((16, 19, 4), np.uint8)}
            for order, key in [(1, "ellipse"), (3, "heart")]
        ]

    monkeypatch.setattr(service, "generate_grow_results_swap", generate)
    image = cv2.imencode(".png", np.zeros((16, 19, 3), np.uint8))[1].tobytes()
    form = {"gender": "female", "hair_style": "1,3", "hairline_height": "-1"}
    if densities:
        form["hair_density"] = densities
    client = TestClient(worker.app, client=("127.0.0.1", 50000))
    response = client.post("/api/v1/hair/grow-22", data=form,
                           files={"image_file": ("input.png", image, "image/png")})
    if not densities:
        assert response.status_code == 422
        assert prompts == []
        return
    assert response.status_code == 200
    assert prompts == [f"填充遮罩区域的头发，发量{label}，人物增加一点磨皮" for label in labels]
    results = response.json()["data"]["results"]
    assert [(item["hair_density"], item["order"]) for item in results] == [
        (density, order) for density in densities for order in (1, 3)
    ]
    for item in results:
        jpg = np.frombuffer(base64.b64decode(item["grown_image_base64"]), np.uint8)
        assert cv2.imdecode(jpg, cv2.IMREAD_COLOR).shape[:2] == (16, 19)


@pytest.mark.parametrize("dy", [-5, 5])
def test_height_moves_paste_and_redraw_band_after_swap(monkeypatch, dy):
    from face_analysis import hairline_grow as grow

    image = np.zeros((80, 80, 3), np.uint8)
    source_mask = np.zeros((80, 80), dtype=bool)
    source_mask[25:45, 20:60] = True
    inner = np.array([[20, 42], [59, 42]])
    outer = np.array([[20, 32], [59, 32]])
    upper = np.ones_like(source_mask)
    monkeypatch.setattr(grow.detector, "detect", lambda image: object())
    monkeypatch.setattr(grow, "estimate_scale_factor", lambda *args: 50)
    monkeypatch.setattr(grow, "compute_mask", lambda *args, **kwargs: (
        source_mask, {"_inner_pts": inner, "_outer_pts": outer, "_upper_mask": upper,
                      "mask_pixels": int(source_mask.sum())}))

    def swap(original, hairline_id, is_hr, ext_mask, denoising_strength, **kwargs):
        assert original is image
        assert ext_mask is source_mask
        generated = np.zeros_like(image)
        generated[30, :, :] = 180
        return generated

    def composite(original, generated, mask, *args, **kwargs):
        assert np.flatnonzero(mask[:, 40]).tolist() == list(range(25 + dy, 45 + dy))
        assert np.flatnonzero(generated[:, 40, 0]).tolist() == [30 + dy]
        return np.where(mask[..., None], generated, original), mask.astype(np.float32)

    def encode_red_mask(mask, h, w):
        assert np.flatnonzero(mask[:, 40]).tolist() == list(range(27 + dy, 38 + dy))
        return "verified"

    monkeypatch.setattr(grow, "_call_swap", swap)
    monkeypatch.setattr(grow, "_composite", composite)
    monkeypatch.setattr(grow, "_red_mask_b64", encode_red_mask)
    result = grow.generate_hairline_redraw(image, "chang_tuoyuan", paste_dy=dy)
    assert result["steps"]["redraw_band_mask_base64"] == "verified"