"""In-process tests that exercise the server's auth middleware via ASGI.

Uses ``httpx.AsyncClient`` with ``ASGITransport`` so no uvicorn / network is
required. Validates that a request forged with the real token-minting
pipeline reaches the FastMCP layer, and that tampering with any of the
pieces is rejected with 4xx.
"""

from __future__ import annotations

import json

import httpx
import pytest

from poc.http_sig import sign_request
from poc.server import build_app
from poc.tokens import mint_ect, mint_mandate


pytestmark = pytest.mark.asyncio


def _headers_for(
    identities,
    *,
    body: bytes,
    audience: str,
    exec_act: str,
    tamper_body: bool = False,
    tamper_aud: bool = False,
) -> tuple[dict[str, str], bytes]:
    """Build a full set of ACT+ECT+signature headers for one request."""
    agent = identities["agent"]
    user = identities["user"]
    mandate = mint_mandate(
        user=user, agent=agent, audience=audience, purpose="test"
    )
    ect = mint_ect(
        agent=agent,
        audience=audience,
        exec_act=exec_act,
        pred_jtis=[mandate.mandate.jti],
        inp_body=body,
    )
    sign_body = b"tampered" if tamper_body else body
    sign_aud = "wrong-audience" if tamper_aud else audience
    signed = sign_request(
        method="POST",
        target_uri="http://testserver/mcp",
        body=sign_body,
        wimse_ect=ect.compact,
        wimse_aud=sign_aud,
        keyid=agent.kid,
        private_key=agent.private_key,
    )
    headers = {
        "content-type": "application/json",
        "accept": "application/json, text/event-stream",
        "authorization": f"Bearer {mandate.compact}",
        "wimse-ect": ect.compact,
        "content-digest": signed.content_digest,
        "signature-input": signed.signature_input,
        "signature": signed.signature,
    }
    return headers, body


from contextlib import asynccontextmanager


@asynccontextmanager
async def _client_for(identities):
    """Return an httpx client wired to the ASGI app with lifespan started.

    FastMCP's streamable-HTTP transport allocates a task group during
    ``lifespan.startup``; ``ASGITransport`` does not run lifespan by
    default, so we manage it explicitly here.
    """
    app = build_app(identities)
    async with app.router.lifespan_context(app):
        transport = httpx.ASGITransport(app=app)
        async with httpx.AsyncClient(
            transport=transport, base_url="http://testserver"
        ) as client:
            yield client


async def test_no_auth_headers_returns_401(identities):
    async with _client_for(identities) as c:
        r = await c.post("/mcp", content=b'{"jsonrpc":"2.0","method":"initialize"}')
        assert r.status_code == 401
        assert "Authorization" in r.text


async def test_valid_initialize_request_is_accepted(identities, tmp_path, monkeypatch):
    monkeypatch.setenv("POC_AUDIT_LOG", str(tmp_path / "audit.jsonl"))
    body = json.dumps({
        "jsonrpc": "2.0",
        "id": 1,
        "method": "initialize",
        "params": {
            "protocolVersion": "2025-03-26",
            "capabilities": {},
            "clientInfo": {"name": "poc-test", "version": "0"},
        },
    }).encode("utf-8")
    headers, _ = _headers_for(
        identities, body=body, audience="mcp-server",
        exec_act="mcp.session.initialize",
    )
    async with _client_for(identities) as c:
        r = await c.post("/mcp", content=body, headers=headers)
        # FastMCP may respond 200 with a session id, or 202, or stream SSE.
        assert r.status_code < 400, f"unexpected status {r.status_code}: {r.text}"


async def test_tampered_body_is_rejected(identities):
    body = json.dumps({"jsonrpc": "2.0", "id": 1, "method": "initialize"}).encode()
    headers, _ = _headers_for(
        identities, body=body, audience="mcp-server",
        exec_act="mcp.session.initialize", tamper_body=True,
    )
    async with _client_for(identities) as c:
        r = await c.post("/mcp", content=body, headers=headers)
        assert r.status_code == 401
        assert "content-digest" in r.text.lower() or "http-signature" in r.text.lower()


async def test_wrong_wimse_aud_is_rejected(identities):
    body = json.dumps({"jsonrpc": "2.0", "id": 1, "method": "initialize"}).encode()
    headers, _ = _headers_for(
        identities, body=body, audience="mcp-server",
        exec_act="mcp.session.initialize", tamper_aud=True,
    )
    async with _client_for(identities) as c:
        r = await c.post("/mcp", content=body, headers=headers)
        assert r.status_code == 401


async def test_unauthorised_tool_is_rejected(identities):
    """A tools/call whose exec_act is not in mandate.cap → 403."""
    body = json.dumps({
        "jsonrpc": "2.0", "id": 1, "method": "tools/call",
        "params": {"name": "search", "arguments": {"query": "x"}},
    }).encode()
    # Craft a mandate where we strip search out of cap by going through
    # the token API: we fabricate headers with the right exec_act but a
    # mandate whose cap doesn't contain it.
    from poc.keys import Identity  # noqa
    from poc.tokens import MCP_CAPS, mint_mandate
    agent = identities["agent"]
    user = identities["user"]

    # Replace MCP_CAPS monkey-patch-free: build mandate directly.
    from act.token import ACTMandate, Capability, TaskClaim
    from act.crypto import sign as act_sign
    from act.token import encode_jws
    import time, uuid
    iat = int(time.time())
    mandate = ACTMandate(
        alg="ES256", kid=user.kid, iss="user", sub="agent",
        aud="mcp-server", iat=iat, exp=iat + 600, jti=str(uuid.uuid4()),
        wid="agent",
        task=TaskClaim(purpose="p", created_by="user"),
        cap=[Capability(action="mcp.summarize")],  # search MISSING
    )
    mandate.validate()
    mandate_compact = encode_jws(mandate, act_sign(user.private_key, mandate.signing_input()))

    ect = mint_ect(
        agent=agent, audience="mcp-server",
        exec_act="mcp.search", pred_jtis=[mandate.jti], inp_body=body,
    )
    signed = sign_request(
        method="POST", target_uri="http://testserver/mcp",
        body=body, wimse_ect=ect.compact, wimse_aud="mcp-server",
        keyid=agent.kid, private_key=agent.private_key,
    )
    headers = {
        "content-type": "application/json",
        "authorization": f"Bearer {mandate_compact}",
        "wimse-ect": ect.compact,
        "content-digest": signed.content_digest,
        "signature-input": signed.signature_input,
        "signature": signed.signature,
    }
    async with _client_for(identities) as c:
        r = await c.post("/mcp", content=body, headers=headers)
        assert r.status_code == 403
        assert "exec_act" in r.text or "cap" in r.text