"""Gateway integration for interfaces 52 and 22: raw forwarding, token and image URLs."""
import asyncio
import base64
import importlib
import io
from contextlib import closing
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock

import httpx
from PIL import Image
import pytest
from fastapi.testclient import TestClient

from gateway import config, forward

gateway = importlib.import_module("gateway.app")
PATH = "/api/v1/hairline/generate-52"
GROW22 = "/api/v1/hair/grow-22"


@pytest.fixture
def client_headers(monkeypatch):
    monkeypatch.setenv("HAIR_API_TOKEN_SECRET", "gateway-test-signing-secret-at-least-32")
    client = TestClient(gateway.app)
    try:
        response = client.post("/api/v1/auth/login", json={"username": "yonghe", "password": "yonghe1999"})
    finally:
        client.close()
    return {"Authorization": "Bearer " + response.json()["data"]["access_token"]}


def image_bytes(fmt):
    buf = io.BytesIO()
    Image.new("RGB", (6, 8), (127, 127, 127)).save(buf, format=fmt)
    return buf.getvalue()


@pytest.fixture
def proxy_setup(tmp_path, monkeypatch):
    output = tmp_path / "annotations"
    output.mkdir()
    cfg = {"shared_password": "private-gateway-token", "public_base_url": "http://testserver",
           "static_dir": str(output), "dispatch": {"request_timeout_seconds": 600, "max_retries": 0}}
    monkeypatch.setattr(config, "get_config", lambda: cfg)
    acquire = AsyncMock(return_value=True)
    release = Mock()
    monkeypatch.setattr(forward, "acquire_global_slot", acquire)
    monkeypatch.setattr(forward, "release_global_slot", release)
    monkeypatch.setattr(forward, "acquire_worker", AsyncMock(return_value=SimpleNamespace(url="http://worker")))
    release_worker = AsyncMock()
    monkeypatch.setattr(forward, "release_worker", release_worker)
    unhealthy = AsyncMock()
    monkeypatch.setattr(forward, "mark_worker_unhealthy", unhealthy)
    static = next(route.app for route in gateway.app.routes if getattr(route, "name", None) == "static")
    monkeypatch.setattr(static, "all_directories", [str(tmp_path), *static.all_directories])
    return cfg, acquire, release, release_worker, unhealthy


@pytest.mark.parametrize("path", ["/api/v1/hairline/generate", PATH, "/api/v1/face/b-end-points"])
def test_routes_forward_to_their_own_worker_path(monkeypatch, path, client_headers):
    monkeypatch.setattr(config, "get_config", lambda: {})
    proxy = AsyncMock(return_value={"code": 0, "data": None})
    monkeypatch.setattr(gateway, "_proxy", proxy)
    client = TestClient(gateway.app, client=("203.0.113.9", 50000))
    try:
        assert client.post(path, headers=client_headers if path != "/api/v1/hairline/generate" else {}).json() == {"code": 0, "data": None}
    finally:
        client.close()
    proxy.assert_awaited_once()
    assert proxy.await_args.args[1] == path


@pytest.mark.parametrize("generate_grow_image", [False, True])
def test_interface52_raw_forwarding_and_served_image_urls(proxy_setup, monkeypatch, generate_grow_image, client_headers):
    cfg, acquire, release, release_worker, unhealthy = proxy_setup
    png, jpg = image_bytes("PNG"), image_bytes("JPEG")
    metadata = {"coordinate_space": "processed_image", "width": 6, "height": 8,
                "framing": {"status": "reference_crop", "axis_x_error_px": 0},
                "height_rule": {"target_center": {"x": 3.0, "y": 2.0}},
                "guide_landmarks": {"9": {"x": 3.0, "y": 3.0}},
                "best_hairline_center_point": {"x": 3, "y": 2},
                "high_hairline_center_point": None, "low_hairline_center_point": None,
                "face_measure": None}
    worker_json = {"code": 0, "data": {**metadata,
        "input_image_base64": base64.b64encode(jpg).decode(),
        "original_image_base64": base64.b64encode(jpg).decode(),
        "hairline_images": [{"hairline_type": "ellipse", "order": 1,
            "height_placement": {"target_error_px": 0.01},
            "image_middle_base64": base64.b64encode(png).decode(),
            "image_high_base64": None, "image_low_base64": None,
            "grown_image_base64": base64.b64encode(jpg).decode() if generate_grow_image else None}]}}
    calls = []

    def worker(request):
        calls.append(request)
        return httpx.Response(200, json=worker_json)

    upstream = httpx.AsyncClient(transport=httpx.MockTransport(worker))
    monkeypatch.setattr(forward, "get_client", lambda: upstream)
    fields = {"gender": "male", "hair_style": "1,3", "normalize_image": "true",
              "generate_grow_image": str(generate_grow_image).lower(), "use_mask": "false",
              "prompt": "test prompt"}
    source = httpx.Request("POST", "http://testserver" + PATH,
                           files={"image_file": ("face.png", png, "image/png")}, data=fields)
    raw = source.read()
    client = TestClient(gateway.app, client=("203.0.113.9", 50000))
    try:
        result = client.post(PATH, content=raw, headers={
            **client_headers,
            "Content-Type": source.headers["content-type"], "X-Internal-Token": "untrusted-client-token"})
        assert result.status_code == 200
        data = result.json()["data"]
        assert {key: data[key] for key in metadata} == metadata
        item = data["hairline_images"][0]
        assert item["image_high_base64"] is None and item["image_low_base64"] is None
        assert "image_high_url" not in item and "image_low_url" not in item
        assert item["height_placement"] == {"target_error_px": 0.01}
        for key, expected in [("input_image_url", jpg), ("original_image_url", jpg)]:
            assert client.get(data[key]).content == expected
            assert data[key].endswith(".jpg")
        assert client.get(item["image_middle_url"]).content == png
        assert item["image_middle_url"].endswith(".png")
        if generate_grow_image:
            assert client.get(item["grown_image_url"]).content == jpg
        else:
            assert item["grown_image_base64"] is None and "grown_image_url" not in item
    finally:
        client.close()
        asyncio.run(upstream.aclose())
    assert len(calls) == 1
    assert str(calls[0].url) == "http://worker" + PATH
    assert calls[0].content == raw
    assert calls[0].headers["content-type"] == source.headers["content-type"]
    assert calls[0].headers["x-internal-token"] == cfg["shared_password"]
    assert calls[0].headers.get("x-forwarded-for") is None
    assert calls[0].headers.get("authorization") is None
    acquire.assert_awaited_once_with(cfg, 0)
    release.assert_called_once()
    release_worker.assert_awaited_once()
    unhealthy.assert_not_awaited()


def test_worker_ip_rejection_is_preserved(proxy_setup, monkeypatch, client_headers):
    _, _, release, release_worker, unhealthy = proxy_setup
    body = {"code": 1010, "message": "来源IP不在接口白名单", "data": None}
    upstream = httpx.AsyncClient(transport=httpx.MockTransport(lambda request: httpx.Response(403, json=body)))
    monkeypatch.setattr(forward, "get_client", lambda: upstream)
    client = TestClient(gateway.app)
    try:
        response = client.post(PATH, headers=client_headers, data={"gender": "male", "hair_style": "1"})
        assert response.status_code == 403 and response.json() == body
    finally:
        client.close()
        asyncio.run(upstream.aclose())
    release.assert_called_once()
    release_worker.assert_awaited_once()
    unhealthy.assert_not_awaited()


def test_public_interface52_page_has_no_internal_password():
    client = TestClient(gateway.app)
    try:
        page = client.get("/static/test_interface52.html")
        assert page.status_code == 200
        assert PATH in page.text and "X-Internal-Token" not in page.text
        assert 'id="client-auth"' in page.text and "X-API-Token" not in page.text
        assert "yonghe1999" not in page.text
        assert "localhost:8187" not in page.text
        assert "test_interface52_compare.html" not in page.text
        assert client.get("/static/hairline_reference_box.js").status_code == 200
    finally:
        client.close()


def test_interface22_forwards_repeated_density_fields_and_serves_image_urls(
    proxy_setup, monkeypatch, client_headers,
):
    cfg, acquire, release, release_worker, unhealthy = proxy_setup
    png, jpg = image_bytes("PNG"), image_bytes("JPEG")
    calls = []

    def worker(request):
        calls.append(request)
        return httpx.Response(200, json={"code": 0, "data": {"results": [
            {"order": 1, "hairline_type": "ellipse", "hair_density": "sparse",
             "image_base64": base64.b64encode(png).decode(),
             "grown_image_base64": base64.b64encode(jpg).decode()},
            {"order": 1, "hairline_type": "ellipse", "hair_density": "dense",
             "image_base64": base64.b64encode(png).decode(), "grown_image_base64": None},
        ]}})

    upstream = httpx.AsyncClient(transport=httpx.MockTransport(worker))
    monkeypatch.setattr(forward, "get_client", lambda: upstream)
    source = httpx.Request(
        "POST", "http://testserver" + GROW22,
        files={"image_file": ("face.png", png, "image/png")},
        data={"gender": "female", "hair_style": "1", "hairline_height": "-0.5",
              "hair_density": ["sparse", "dense"]},
    )
    raw = source.read()
    assert raw.count(b'name="hair_density"') == 2
    try:
        with closing(TestClient(gateway.app)) as client:
            response = client.post(GROW22, content=raw, headers={
                **client_headers, "Content-Type": source.headers["content-type"],
            })
            assert response.status_code == 200
            sparse, dense = response.json()["data"]["results"]
            assert [sparse["hair_density"], dense["hair_density"]] == ["sparse", "dense"]
            assert client.get(sparse["image_url"]).content == png
            assert client.get(sparse["grown_image_url"]).content == jpg
            assert dense["grown_image_base64"] is None
    finally:
        asyncio.run(upstream.aclose())
    assert len(calls) == 1
    assert str(calls[0].url) == "http://worker" + GROW22
    assert calls[0].content == raw
    assert calls[0].headers["x-internal-token"] == cfg["shared_password"]
    assert "authorization" not in calls[0].headers
    acquire.assert_awaited_once_with(cfg, 0)
    release.assert_called_once()
    release_worker.assert_awaited_once()
    unhealthy.assert_not_awaited()


def test_interface22_page_uses_gateway_login_and_authenticated_requests():
    with closing(TestClient(gateway.app)) as client:
        page = client.get("/static/test_interface22.html")
        assert page.status_code == 200
        assert 'id="client-auth"' in page.text
        for path in (GROW22, PATH):
            assert f"window.hairClientFetch('{path}'" in page.text
        assert "X-Internal-Token" not in page.text
        assert "localhost:8187" not in page.text
        assert client.get("/static/client_auth.js").status_code == 200