"""中継本体。バイト単位の無改変・アドミッション・認証・応答の扱い。"""

from __future__ import annotations

import gzip
import hashlib
import json

import httpx
import pytest

from app import logging_utils

from .conftest import INTERNAL_KEY, AsyncChunks, auth, create_key, stream_response

LLAMACPP_BODY = json.dumps(
    {
        "model": "qwen3-8b",
        "prompt": "日本語のプロンプト ❤ with emoji 🎌",
        "grammar": 'root ::= "yes" | "no"',
        "cache_prompt": True,
        "n_probs": 5,
        "stop": ["</s>"],
        "samplers": ["top_k", "typ_p", "temperature"],
    },
    ensure_ascii=False,
).encode("utf-8")


# ---------------------------------------------------------------- 無改変中継


async def test_request_body_is_forwarded_byte_for_byte(db, keys, settings, build_proxy, recorder):
    secret = await create_key(db, keys, settings)
    client, _ = build_proxy()
    async with client:
        response = await client.post(
            "/v1/chat/completions",
            headers={**auth(secret), "content-type": "application/json"},
            content=LLAMACPP_BODY,
        )

    assert response.status_code == 200
    forwarded = recorder.last.content
    assert forwarded == LLAMACPP_BODY
    assert hashlib.sha256(forwarded).hexdigest() == hashlib.sha256(LLAMACPP_BODY).hexdigest()


async def test_llamacpp_parameters_survive(db, keys, settings, build_proxy, recorder):
    """再シリアライズすると鍵順や数値表現が変わりうる。原本のまま送る。"""

    secret = await create_key(db, keys, settings)
    client, _ = build_proxy()
    quirky = b'{"model":"m","n_probs":5,"cache_prompt":true,"temperature":1.0,"grammar":"root ::= \\"a\\""}'
    async with client:
        await client.post(
            "/v1/completions",
            headers={**auth(secret), "content-type": "application/json"},
            content=quirky,
        )
    assert recorder.last.content == quirky


async def test_path_and_query_are_preserved(db, keys, settings, build_proxy, recorder):
    secret = await create_key(db, keys, settings)
    client, _ = build_proxy()
    async with client:
        await client.get("/v1/some%20path/x?a=1&b=%E3%81%82&flag", headers=auth(secret))

    url = recorder.last.url
    assert url.host == "upstream.test"
    assert url.raw_path == b"/v1/some%20path/x?a=1&b=%E3%81%82&flag"


async def test_response_body_is_not_decompressed(db, keys, settings, build_proxy):
    """``aiter_raw`` で受けるため、上流の圧縮はそのまま素通しする。"""

    secret = await create_key(db, keys, settings)
    payload = gzip.compress(b'{"choices":[]}')

    def respond(request: httpx.Request) -> httpx.Response:
        return httpx.Response(
            200,
            headers=[
                (b"content-type", b"application/json"),
                (b"content-encoding", b"gzip"),
                (b"content-length", str(len(payload)).encode()),
            ],
            stream=AsyncChunks([payload]),
        )

    client, _ = build_proxy(respond)
    async with client:
        async with client.stream(
            "POST",
            "/v1/chat/completions",
            headers={**auth(secret), "accept-encoding": "gzip"},
            content=b'{"model":"m"}',
        ) as response:
            raw = b"".join([chunk async for chunk in response.aiter_raw()])
            assert response.headers["content-encoding"] == "gzip"

    assert raw == payload


async def test_sse_chunks_are_passed_through(db, keys, settings, build_proxy):
    secret = await create_key(db, keys, settings)
    chunks = [
        b'data: {"choices":[{"delta":{"content":"a"}}]}\n\n',
        b'data: {"choices":[{"delta":{"content":"b"}}]}\n\n',
        b"data: [DONE]\n\n",
    ]

    def respond(request: httpx.Request) -> httpx.Response:
        return httpx.Response(
            200,
            headers=[(b"content-type", b"text/event-stream")],
            stream=AsyncChunks(chunks),
        )

    client, _ = build_proxy(respond)
    async with client:
        async with client.stream(
            "POST", "/v1/chat/completions", headers=auth(secret), content=b'{"model":"m","stream":true}'
        ) as response:
            received = [chunk async for chunk in response.aiter_raw()]
            assert response.headers["content-type"] == "text/event-stream"

    assert b"".join(received) == b"".join(chunks)


async def test_streaming_usage_is_recorded(db, keys, settings, build_proxy, logs):
    """クライアントが usage を要求した場合のみ記録できる。"""

    secret = await create_key(db, keys, settings)
    chunks = [
        b'data: {"choices":[{"delta":{"content":"a"}}]}\n\n',
        b'data: {"choices":[],"usage":{"prompt_tokens":11,"completion_tokens":22}}\n\n',
        b"data: [DONE]\n\n",
    ]

    def respond(request: httpx.Request) -> httpx.Response:
        return httpx.Response(
            200, headers=[(b"content-type", b"text/event-stream")], stream=AsyncChunks(chunks)
        )

    client, _ = build_proxy(respond)
    async with client:
        await client.post(
            "/v1/chat/completions",
            headers=auth(secret),
            content=b'{"model":"m","stream":true,"stream_options":{"include_usage":true}}',
        )
    await logs.stop()

    row = await db.fetchone("SELECT prompt_tokens, completion_tokens FROM request_logs ORDER BY id DESC LIMIT 1")
    assert row is not None
    assert (row["prompt_tokens"], row["completion_tokens"]) == (11, 22)


async def test_non_streaming_usage_is_recorded(db, keys, settings, build_proxy, logs):
    secret = await create_key(db, keys, settings)
    body = json.dumps({"choices": [], "usage": {"prompt_tokens": 3, "completion_tokens": 4}}).encode()

    client, _ = build_proxy(lambda request: stream_response(200, content=body))
    async with client:
        await client.post("/v1/chat/completions", headers=auth(secret), content=b'{"model":"m"}')
    await logs.stop()

    row = await db.fetchone("SELECT prompt_tokens, completion_tokens FROM request_logs ORDER BY id DESC LIMIT 1")
    assert (row["prompt_tokens"], row["completion_tokens"]) == (3, 4)


# ------------------------------------------------------------------ 認証


async def test_missing_bearer_is_401(db, keys, settings, build_proxy, recorder):
    client, _ = build_proxy()
    async with client:
        response = await client.post("/v1/chat/completions", content=b"{}")
    assert response.status_code == 401
    assert response.json()["error"]["code"] == "invalid_api_key"
    assert recorder.requests == []


async def test_unknown_key_is_401(db, keys, settings, build_proxy):
    client, _ = build_proxy()
    async with client:
        response = await client.post(
            "/v1/chat/completions", headers=auth("sk-aig-" + "z" * 32), content=b"{}"
        )
    assert response.status_code == 401


async def test_disabled_key_is_401(db, keys, settings, build_proxy):
    secret = await create_key(db, keys, settings, enabled=False)
    client, _ = build_proxy()
    async with client:
        response = await client.post("/v1/chat/completions", headers=auth(secret), content=b"{}")
    assert response.status_code == 401


async def test_expired_key_is_401(db, keys, settings, build_proxy):
    secret = await create_key(db, keys, settings, expires_at="2000-01-01T00:00:00+09:00")
    client, _ = build_proxy()
    async with client:
        response = await client.post("/v1/chat/completions", headers=auth(secret), content=b"{}")
    assert response.status_code == 401


# ------------------------------------------------------ アドミッション コントロール


async def test_oversized_body_is_413(db, keys, settings, build_proxy, recorder):
    secret = await create_key(db, keys, settings)
    client, effective = build_proxy()
    body = b"x" * (effective.max_request_body_bytes + 1)
    async with client:
        response = await client.post("/v1/chat/completions", headers=auth(secret), content=body)
    assert response.status_code == 413
    assert response.json()["error"]["code"] == "request_too_large"
    assert recorder.requests == []


async def test_oversized_grammar_is_400(db, keys, settings, build_proxy, recorder):
    secret = await create_key(db, keys, settings)
    client, effective = build_proxy()
    body = json.dumps({"model": "m", "grammar": "g" * (effective.max_grammar_bytes + 1)}).encode()
    async with client:
        response = await client.post(
            "/v1/chat/completions",
            headers={**auth(secret), "content-type": "application/json"},
            content=body,
        )
    assert response.status_code == 400
    assert response.json()["error"]["code"] == "grammar_too_large"
    assert recorder.requests == []


async def test_excessive_n_probs_is_400(db, keys, settings, build_proxy, recorder):
    secret = await create_key(db, keys, settings)
    client, effective = build_proxy()
    body = json.dumps({"model": "m", "n_probs": effective.max_n_probs + 1}).encode()
    async with client:
        response = await client.post(
            "/v1/chat/completions",
            headers={**auth(secret), "content-type": "application/json"},
            content=body,
        )
    assert response.status_code == 400
    assert recorder.requests == []


# ------------------------------------------------------------------ パス方針


@pytest.mark.parametrize("path", ["/slots", "/slots/0", "/props", "/metrics", "/lora-adapters"])
async def test_denied_paths_are_404(db, keys, settings, build_proxy, recorder, path):
    secret = await create_key(db, keys, settings)
    client, _ = build_proxy()
    async with client:
        response = await client.get(path, headers=auth(secret))
    assert response.status_code == 404
    assert recorder.requests == [], "拒否パスを上流へ転送しています"


async def test_denied_path_is_404_even_without_a_key(db, keys, settings, build_proxy, recorder):
    """存在の有無も漏らさない。認証より先に拒否する。"""

    client, _ = build_proxy()
    async with client:
        response = await client.get("/slots", headers={})
    assert response.status_code == 404


async def test_similar_paths_are_not_denied(db, keys, settings, build_proxy, recorder):
    secret = await create_key(db, keys, settings)
    client, _ = build_proxy()
    async with client:
        response = await client.get("/slotsfoo", headers=auth(secret))
    assert response.status_code == 200
    assert len(recorder.requests) == 1


async def test_allowlist_mode_rejects_unlisted(db, keys, settings, build_proxy, recorder):
    secret = await create_key(db, keys, settings)
    client, _ = build_proxy(overrides={"path_policy": "allowlist"})
    async with client:
        denied = await client.get("/health", headers=auth(secret))
        allowed = await client.get("/v1/models", headers=auth(secret))
    assert denied.status_code == 404
    assert allowed.status_code == 200


# ------------------------------------------------------------ /v1/models 絞り込み


async def test_models_response_is_filtered(db, keys, settings, build_proxy):
    secret = await create_key(db, keys, settings, models=["qwen3-8b"])
    payload = json.dumps(
        {"object": "list", "data": [{"id": "qwen3-8b"}, {"id": "gemma3-12b"}]}
    ).encode()

    client, _ = build_proxy(lambda request: stream_response(200, content=payload))
    async with client:
        response = await client.get("/v1/models", headers=auth(secret))

    assert response.status_code == 200
    assert [item["id"] for item in response.json()["data"]] == ["qwen3-8b"]
    assert int(response.headers["content-length"]) == len(response.content)


async def test_models_response_is_untouched_for_unrestricted_key(db, keys, settings, build_proxy):
    secret = await create_key(db, keys, settings, models=[])
    payload = json.dumps({"object": "list", "data": [{"id": "a"}, {"id": "b"}]}).encode()

    client, _ = build_proxy(lambda request: stream_response(200, content=payload))
    async with client:
        response = await client.get("/v1/models", headers=auth(secret))

    assert response.content == payload


async def test_models_filter_can_be_disabled(db, keys, settings, build_proxy):
    secret = await create_key(db, keys, settings, models=["a"])
    payload = json.dumps({"object": "list", "data": [{"id": "a"}, {"id": "b"}]}).encode()

    client, _ = build_proxy(
        lambda request: stream_response(200, content=payload),
        overrides={"filter_models_response": False},
    )
    async with client:
        response = await client.get("/v1/models", headers=auth(secret))

    assert response.content == payload


# ------------------------------------------------------------ 内部キーの安全網


async def test_internal_key_in_error_body_is_redacted(db, keys, settings, build_proxy):
    """上流が誤って内部キーを返しても、クライアントへは出さない。"""

    secret = await create_key(db, keys, settings)
    leaking = json.dumps({"error": f"invalid api key: {INTERNAL_KEY}"}).encode()

    client, _ = build_proxy(lambda request: stream_response(401, content=leaking))
    async with client:
        response = await client.post(
            "/v1/chat/completions", headers=auth(secret), content=b'{"model":"m"}'
        )

    assert response.status_code == 401
    assert INTERNAL_KEY.encode() not in response.content
    assert logging_utils.REDACTED.encode() in response.content
    assert int(response.headers["content-length"]) == len(response.content)


async def test_clean_error_body_is_passed_through_unchanged(db, keys, settings, build_proxy):
    secret = await create_key(db, keys, settings)
    body = json.dumps({"error": {"message": "bad request"}}).encode()

    client, _ = build_proxy(lambda request: stream_response(400, content=body))
    async with client:
        response = await client.post(
            "/v1/chat/completions", headers=auth(secret), content=b'{"model":"m"}'
        )

    assert response.status_code == 400
    assert response.content == body


async def test_internal_key_is_not_in_response_headers(db, keys, settings, build_proxy):
    secret = await create_key(db, keys, settings)
    client, _ = build_proxy()
    async with client:
        response = await client.post(
            "/v1/chat/completions", headers=auth(secret), content=b'{"model":"m"}'
        )

    joined = " ".join(f"{name}: {value}" for name, value in response.headers.items())
    assert INTERNAL_KEY not in joined
    assert secret not in joined


async def test_large_error_body_is_streamed_without_truncation(db, keys, settings, build_proxy):
    """走査上限を超えるエラー本文も欠けずに届く。"""

    secret = await create_key(db, keys, settings)
    body = b"E" * (128 * 1024)

    client, _ = build_proxy(lambda request: stream_response(500, content=body))
    async with client:
        response = await client.post(
            "/v1/chat/completions", headers=auth(secret), content=b'{"model":"m"}'
        )

    assert response.status_code == 500
    assert response.content == body


# ------------------------------------------------------------------ 上流障害


async def test_upstream_connect_error_is_502(db, keys, settings, build_proxy):
    def boom(request: httpx.Request) -> httpx.Response:
        raise httpx.ConnectError("refused", request=request)

    secret = await create_key(db, keys, settings)
    client, _ = build_proxy(boom)
    async with client:
        response = await client.post(
            "/v1/chat/completions", headers=auth(secret), content=b'{"model":"m"}'
        )

    assert response.status_code == 502
    assert response.json()["error"]["type"] == "api_error"


async def test_upstream_read_timeout_is_504(db, keys, settings, build_proxy):
    def boom(request: httpx.Request) -> httpx.Response:
        raise httpx.ReadTimeout("idle", request=request)

    secret = await create_key(db, keys, settings)
    client, _ = build_proxy(boom)
    async with client:
        response = await client.post(
            "/v1/chat/completions", headers=auth(secret), content=b'{"model":"m"}'
        )

    assert response.status_code == 504
    assert response.json()["error"]["code"] == "upstream_timeout"