"""Gateway-only interface 52 integration: raw forwarding, token and image URLs."""
import asyncio
import base64
import importlib
import io
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"
CLIENT_HEADERS = {"X-API-Token": "yonghe1999"}


@pytest.mark.parametrize("headers", [
    {}, {"X-API-Token": ""}, {"X-API-Token": "wrong"},
    {"X-API-Token": "yonghe1999 "}, {"X-Internal-Token": "yonghe1999"},
])
def test_invalid_client_token_is_rejected_before_dispatch(monkeypatch, headers):
    monkeypatch.setattr(config, "get_config", lambda: {})
    proxy = AsyncMock()
    monkeypatch.setattr(gateway, "_proxy", proxy)
    client = TestClient(gateway.app)
    try:
        response = client.post(PATH, headers=headers)
        assert response.status_code == 401
        body = response.json()
        assert body["code"] == 1009 and body["data"] is None
        assert body["request_id"].startswith("gw-")
    finally:
        client.close()
    proxy.assert_not_awaited()


def test_custom_token_replaces_default_only_on_interface52(monkeypatch):
    monkeypatch.setattr(config, "get_config", lambda: {"interface52_token": "custom-client-token"})
    proxy = AsyncMock(return_value={"code": 0, "data": None})
    monkeypatch.setattr(gateway, "_proxy", proxy)
    client = TestClient(gateway.app)
    try:
        assert client.post(PATH, headers=CLIENT_HEADERS).status_code == 401
        assert client.post(PATH, headers={"X-API-Token": "custom-client-token"}).status_code == 200
        assert client.post("/api/v1/hairline/generate").status_code == 200
    finally:
        client.close()
    assert proxy.await_count == 2


@pytest.mark.parametrize("token", ["", None, 123, []])
def test_invalid_configured_token_cannot_disable_auth(monkeypatch, token):
    monkeypatch.setattr(config, "get_config", lambda: {"interface52_token": token})
    proxy = AsyncMock()
    monkeypatch.setattr(gateway, "_proxy", proxy)
    client = TestClient(gateway.app)
    try:
        assert client.post(PATH, headers=CLIENT_HEADERS).status_code == 401
    finally:
        client.close()
    proxy.assert_not_awaited()


@pytest.mark.parametrize("token", ["", None, 123, []])
def test_startup_rejects_invalid_client_token_configuration(tmp_path, monkeypatch, token):
    import json
    path = tmp_path / "config.json"
    path.write_text(json.dumps({"interface52_token": token}))
    monkeypatch.setattr(config, "CONFIG_PATH", path)
    monkeypatch.setattr(config, "_config_cache", None)
    with pytest.raises(ValueError, match="interface52_token"):
        config.load_config()


def test_existing_config_gets_default_client_token(tmp_path, monkeypatch):
    import json
    path = tmp_path / "config.json"
    path.write_text(json.dumps({"workers": ["http://worker"], "shared_password": "test-internal-token",
                                "static_dir": str(tmp_path)}))
    monkeypatch.setattr(config, "CONFIG_PATH", path)
    monkeypatch.setattr(config, "_config_cache", None)
    assert config.load_config()["interface52_token"] == "yonghe1999"


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])
def test_hairline_routes_forward_to_their_own_worker_path(monkeypatch, path):
    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 == PATH 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):
    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("x-api-token") 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):
    _, _, 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="apiToken"' in page.text and "X-API-Token" 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()