"""JWT auth enforcement tests against the deployed QA routers.

These tests do not exercise GraphQL queries — every request is a trivial
`{ __typename }`. What varies is the Authorization header, and what is
asserted is each router variant's pass/block decision:

- public (config-qa.yaml): every auth rule is `Warning` — invalid auth is
  logged but the request passes.
- internal (config-qa-internal.yaml) and mcp (config-qa-mcp.yaml): every
  auth rule is `Block` — invalid auth gets the plugin's uniform 401
  (`Unauthorized.` / `UNAUTHORIZED`).

Setting LOCAL_ROUTER_URL replaces the matrix with that single local router
(LOCAL_ROUTER_MODE=warn|block, default block); the
`make test_integration_local_<public|internal|mcp>` targets do this against
the matching docker compose router.

Optional env vars for credentialed cases (skipped when unset):
- EXPIRED_JWT: a real expired token (cannot be minted without the IdP key).
"""

import base64
import json
import os
from typing import Any

import httpx
import pytest

QUERY = '{ __typename }'

# name -> (url, mode); mode is 'warn' (passes invalid auth) or 'block' (401s).
ROUTERS: dict[str, tuple[str, str]] = {
    'public': ('https://qa-graphql-router.theorchard.io/graphql', 'warn'),
    'internal': ('https://qa-graphql-router-internal.theorchard.io/graphql', 'block'),
    'mcp': ('https://qa-graphql-router-mcp.theorchard.io/graphql', 'block'),
}

if os.environ.get('LOCAL_ROUTER_URL'):
    # A local router replaces the QA matrix entirely — local runs are local only.
    ROUTERS = {
        'local': (
            os.environ['LOCAL_ROUTER_URL'],
            os.environ.get('LOCAL_ROUTER_MODE', 'block'),
        )
    }

ROUTER_PARAMS = [pytest.param(name, id=name) for name in ROUTERS]


def post_graphql(url: str, authorization: str | list[str] | None) -> httpx.Response:
    """POST a trivial GraphQL query with the given Authorization header(s).

    `authorization` may be None (header omitted), a single value, or a list
    of values (sent as repeated Authorization headers).
    """
    headers: list[tuple[str, str]] = [
        ('content-type', 'application/json'),
        # Required by theorchard.require_apollo_client_name.
        ('apollographql-client-name', 'graphql-router-tests'),
    ]
    if isinstance(authorization, str):
        headers.append(('authorization', authorization))
    elif isinstance(authorization, list):
        headers.extend(('authorization', value) for value in authorization)

    return httpx.post(url, content=json.dumps({'query': QUERY}), headers=headers, timeout=30)


def assert_blocked(response: httpx.Response) -> None:
    """Assert the auth enforcement plugin's uniform block response."""
    assert response.status_code == 401, f'expected 401, got {response.status_code}: {response.text}'
    body = response.json()
    assert body['errors'][0]['message'] == 'Unauthorized.'
    assert body['errors'][0]['extensions']['code'] == 'UNAUTHORIZED'


def assert_passed(response: httpx.Response) -> None:
    """Assert the request was not blocked and the query resolved."""
    assert response.status_code == 200, f'expected 200, got {response.status_code}: {response.text}'
    body = response.json()
    assert 'errors' not in body or not body['errors'], f'unexpected errors: {body}'
    assert body['data'] == {'__typename': 'Query'}


def assert_enforced(router: str, authorization: str | list[str] | None) -> None:
    """Assert the router's mode-appropriate reaction to invalid auth."""
    url, mode = ROUTERS[router]
    response = post_graphql(url, authorization)
    if mode == 'block':
        assert_blocked(response)
    else:
        assert_passed(response)


def _b64url(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b'=').decode()


def forge_jwt(kid: str | None) -> str:
    """Build a structurally valid JWT with a junk signature.

    Signature verification can never succeed (we hold no trusted key); with
    an unknown/missing kid the router fails key lookup before the signature
    is ever checked.
    """
    header: dict[str, Any] = {'alg': 'RS256', 'typ': 'JWT'}
    if kid is not None:
        header['kid'] = kid
    payload = {
        'iss': 'https://forged.example.com/',
        'sub': 'forged|user',
        'aud': 'https://qa-ows.theorchard.io',
        'exp': 9_999_999_999,
        'iat': 0,
        'azp': 'forged-client',
    }
    return '.'.join(
        [
            _b64url(json.dumps(header).encode()),
            _b64url(json.dumps(payload).encode()),
            _b64url(b'forged-signature'),
        ]
    )


@pytest.mark.parametrize('router', ROUTER_PARAMS)
class TestJwtEnforcement:
    """Invalid Authorization variations: block-mode 401s, warn-mode passes."""

    def test_missing_authorization_header(self, router: str):
        assert_enforced(router, None)

    def test_bearer_token_that_is_not_a_jwt(self, router: str):
        assert_enforced(router, 'Bearer not-a-jwt')

    def test_jwt_without_kid(self, router: str):
        assert_enforced(router, f'Bearer {forge_jwt(kid=None)}')

    def test_jwt_with_unknown_kid(self, router: str):
        assert_enforced(router, f'Bearer {forge_jwt(kid="unknown-kid")}')

    def test_unknown_auth_scheme(self, router: str):
        assert_enforced(router, 'Basic dXNlcjpwYXNz')

    def test_lowercase_bearer_scheme(self, router: str):
        # Scheme matching is case-sensitive; a lowercase scheme is not
        # treated as a Bearer token (unknown auth).
        assert_enforced(router, 'bearer not-a-jwt')

    def test_multiple_authorization_headers(self, router: str):
        # The load balancer in front of the deployed routers rejects
        # duplicate Authorization headers with its own HTML 400 before the
        # router sees them; only a bare router (local) reaches the
        # multiple_auth rule. Either way the request must not pass as
        # authenticated on block-mode routers.
        url, mode = ROUTERS[router]
        response = post_graphql(url, ['Bearer one', 'Bearer two'])
        if response.status_code == 400:
            return
        if mode == 'block':
            assert_blocked(response)
        else:
            assert_passed(response)

    @pytest.mark.skipif(
        not os.environ.get('EXPIRED_JWT'),
        reason='Set EXPIRED_JWT to a real expired token; one cannot be minted without the IdP key',
    )
    def test_expired_jwt(self, router: str):
        assert_enforced(router, f'Bearer {os.environ["EXPIRED_JWT"]}')


@pytest.mark.parametrize('router', ROUTER_PARAMS)
class TestValidToken:
    """A valid token must pass everywhere; a tampered one must not."""

    def test_valid_token_passes(self, router: str, valid_jwt: str):
        url, _ = ROUTERS[router]
        assert_passed(post_graphql(url, f'Bearer {valid_jwt}'))

    def test_tampered_signature_is_enforced(self, router: str, valid_jwt: str):
        # Real header (known kid) + real claims, broken signature: the one
        # case that exercises actual signature verification.
        suffix = 'AAAA' if not valid_jwt.endswith('AAAA') else 'BBBB'
        assert_enforced(router, f'Bearer {valid_jwt[:-4]}{suffix}')