"""MCP OAuth 2.1 connector flow — discovery, DCR, PKCE, token, MCP use.

Exercises the contract a Claude / Cursor custom connector relies on:

* RFC 9728 / RFC 8414 discovery documents are served and well-formed.
* An unauthenticated MCP probe returns 401 with a ``WWW-Authenticate``
  header pointing at the protected-resource metadata (the bootstrap
  signal a connector follows).
* Dynamic client registration (RFC 7591) issues a public client_id and
  rejects a non-https / non-loopback redirect_uri.
* The full PKCE authorization-code grant: register → consent decision →
  token exchange → the minted access token authenticates a real MCP
  ``tools/call whoami``.
* Refresh-token rotation (single-use), tampered PKCE verifier, replayed
  code, and a wrong-audience token are all rejected.
"""
from __future__ import annotations

import base64
import hashlib
import json
import re
import secrets
import uuid
from contextlib import asynccontextmanager
from datetime import UTC, datetime, timedelta
from urllib.parse import parse_qs, urlparse

import pytest
from sqlalchemy import select, text

from app.auth.models import AuthUserRole
from app.core.security import create_access_token, create_mcp_access_token
from app.identity.person_models import IdentityPerson

_MCP_HEADERS = {
    "Content-Type": "application/json",
    "Accept": "application/json, text/event-stream",
}


async def _make_employee(db_session) -> IdentityPerson:
    user_id = uuid.uuid4()
    user = IdentityPerson(
        id=user_id,
        work_email=f"oauth-{user_id}@example.com",
        display_name="OAuth User",
        employment_status="active",
        source_system="manual",
        hris_person_id=f"manual-{user_id}",
    )
    db_session.add(user)
    await db_session.flush()
    role_id_row = await db_session.execute(
        text("SELECT id FROM auth_roles WHERE name = 'employee'")
    )
    role_id = role_id_row.scalar_one()
    if isinstance(role_id, str):
        role_id = uuid.UUID(role_id)
    db_session.add(AuthUserRole(user_id=user_id, role_id=role_id))
    await db_session.flush()
    return user


def _pkce_pair() -> tuple[str, str]:
    verifier = secrets.token_urlsafe(64)[:96]
    challenge = base64.urlsafe_b64encode(
        hashlib.sha256(verifier.encode("ascii")).digest()
    ).decode("ascii").rstrip("=")
    return verifier, challenge


def _parse_sse(body: str) -> dict:
    m = re.search(r"^data:\s*(\{.*\})\s*$", body, re.MULTILINE)
    assert m, f"no SSE data line: {body!r}"
    return json.loads(m.group(1))


async def _register_client(client, redirect_uri="https://app.example.com/cb") -> str:
    r = await client.post(
        "/api/v1/auth/oauth/register",
        json={
            "client_name": "Test Connector",
            "redirect_uris": [redirect_uri],
            "grant_types": ["authorization_code", "refresh_token"],
            "response_types": ["code"],
            "token_endpoint_auth_method": "none",
        },
    )
    assert r.status_code == 201, r.text
    return r.json()["client_id"]


# ── Discovery ───────────────────────────────────────────────────────

@pytest.mark.asyncio
async def test_protected_resource_metadata(client):
    r = await client.get("/.well-known/oauth-protected-resource/api/v1/mcp")
    assert r.status_code == 200, r.text
    body = r.json()
    assert body["resource"].endswith("/api/v1/mcp")
    assert body["authorization_servers"], "must advertise ≥1 auth server"
    assert "mcp" in body["scopes_supported"]


@pytest.mark.asyncio
async def test_authorization_server_metadata(client):
    r = await client.get("/.well-known/oauth-authorization-server")
    assert r.status_code == 200, r.text
    body = r.json()
    assert body["token_endpoint"].endswith("/api/v1/auth/oauth/token")
    assert body["registration_endpoint"].endswith("/api/v1/auth/oauth/register")
    assert body["code_challenge_methods_supported"] == ["S256"]
    assert "authorization_code" in body["grant_types_supported"]


@pytest.mark.asyncio
async def test_unauthenticated_mcp_probe_advertises_oauth(client):
    r = await client.post("/api/v1/mcp/", headers=_MCP_HEADERS, json={
        "jsonrpc": "2.0", "id": 1, "method": "tools/list", "params": {},
    })
    assert r.status_code == 401
    www = r.headers.get("www-authenticate", "")
    assert "resource_metadata=" in www
    assert "oauth-protected-resource" in www


# ── Dynamic client registration ─────────────────────────────────────

@pytest.mark.asyncio
async def test_register_issues_public_client(client):
    client_id = await _register_client(client)
    assert client_id.startswith("helix_mcp_")


@pytest.mark.asyncio
async def test_register_rejects_bad_redirect(client):
    r = await client.post("/api/v1/auth/oauth/register", json={
        "client_name": "Bad", "redirect_uris": ["http://evil.example.com/cb"],
    })
    assert r.status_code == 400
    assert r.json()["error"] == "invalid_redirect_uri"


@pytest.mark.asyncio
async def test_register_bounds_redirect_metadata(client):
    from app.auth.oauth import service as oauth_service

    too_many = [f"https://app.example.com/cb/{i}" for i in range(
        oauth_service.MAX_REDIRECT_URIS + 1,
    )]
    r = await client.post("/api/v1/auth/oauth/register", json={
        "redirect_uris": too_many,
    })
    assert r.status_code == 400
    assert r.json()["error"] == "invalid_redirect_uri"

    r = await client.post("/api/v1/auth/oauth/register", json={
        "redirect_uris": [
            "https://app.example.com/" + "x" * oauth_service.MAX_REDIRECT_URI_LENGTH
        ],
    })
    assert r.status_code == 400
    assert "exceeds" in r.json()["error_description"]


@pytest.mark.asyncio
async def test_register_rejects_unsupported_protocol_metadata(client):
    r = await client.post("/api/v1/auth/oauth/register", json={
        "redirect_uris": ["https://app.example.com/cb"],
        "grant_types": ["client_credentials"],
    })
    assert r.status_code == 400
    assert r.json()["error"] == "invalid_client_metadata"


@pytest.mark.asyncio
async def test_registration_quota_counts_only_ungranted_clients(
    db_session, monkeypatch,
):
    from app.auth.oauth import service as oauth_service
    from app.auth.oauth.models import OAuthClient, OAuthGrant

    monkeypatch.setattr(oauth_service, "MAX_UNGRANTED_CLIENTS", 100)
    monkeypatch.setattr(oauth_service, "MAX_UNGRANTED_CLIENTS_PER_IP", 1)
    first = OAuthClient(
        id=uuid.uuid4(), client_id="helix_mcp_quota_1", client_name="quota",
        redirect_uris=["https://example.com/cb"], created_ip="203.0.113.1",
    )
    db_session.add(first)
    await db_session.flush()
    with pytest.raises(oauth_service.OAuthError) as per_ip_exc:
        await oauth_service.enforce_registration_quota(
            db_session, created_ip="203.0.113.1",
        )
    assert per_ip_exc.value.status == 429

    # A different source still has headroom until the global bound is hit.
    await oauth_service.enforce_registration_quota(
        db_session, created_ip="203.0.113.2",
    )
    monkeypatch.setattr(oauth_service, "MAX_UNGRANTED_CLIENTS", 1)
    with pytest.raises(oauth_service.OAuthError) as exc:
        await oauth_service.enforce_registration_quota(
            db_session, created_ip="203.0.113.2",
        )
    assert exc.value.status == 503

    # Once a user has granted this client it is durable governed state, not
    # anonymous registration capacity.
    db_session.add(OAuthGrant(
        id=uuid.uuid4(), user_id=uuid.uuid4(), client_id=first.client_id,
    ))
    await db_session.flush()
    await oauth_service.enforce_registration_quota(
        db_session, created_ip="203.0.113.2",
    )


@pytest.mark.asyncio
async def test_registration_cleanup_expires_only_ungranted_clients(db_session):
    from app.auth.oauth import service as oauth_service
    from app.auth.oauth.models import OAuthClient, OAuthGrant

    now = datetime.now(UTC)
    stale_ungranted = OAuthClient(
        id=uuid.uuid4(), client_id="helix_mcp_stale_ungranted",
        redirect_uris=["https://example.com/a"], created_ip="203.0.113.1",
        created_at=now - timedelta(days=8),
    )
    stale_granted = OAuthClient(
        id=uuid.uuid4(), client_id="helix_mcp_stale_granted",
        redirect_uris=["https://example.com/b"], created_ip="203.0.113.1",
        created_at=now - timedelta(days=8),
    )
    db_session.add_all([stale_ungranted, stale_granted])
    await db_session.flush()
    db_session.add(OAuthGrant(
        id=uuid.uuid4(), user_id=uuid.uuid4(),
        client_id=stale_granted.client_id,
    ))
    await db_session.flush()

    assert await oauth_service.prune_stale_ungranted_clients(
        db_session, now=now,
    ) == 1
    rows = (await db_session.execute(select(OAuthClient.client_id))).scalars().all()
    assert stale_ungranted.client_id not in rows
    assert stale_granted.client_id in rows


@pytest.mark.asyncio
async def test_register_fails_closed_when_global_registration_lock_is_busy(
    client, monkeypatch,
):
    from app.auth.oauth import router as oauth_router
    from app.core.distributed_lock import LockBusy

    @asynccontextmanager
    async def _busy(*_args, **_kwargs):
        raise LockBusy("oauth_dynamic_client_registration")
        yield  # pragma: no cover

    monkeypatch.setattr(oauth_router, "distributed_lock", _busy)
    r = await client.post("/api/v1/auth/oauth/register", json={
        "redirect_uris": ["https://app.example.com/cb"],
    })
    assert r.status_code == 429
    assert r.json()["error"] == "temporarily_unavailable"


# ── Authorize endpoint ──────────────────────────────────────────────

@pytest.mark.asyncio
async def test_authorize_unknown_client_400(client):
    _, challenge = _pkce_pair()
    r = await client.get("/api/v1/auth/oauth/authorize", params={
        "response_type": "code", "client_id": "helix_mcp_nope",
        "redirect_uri": "https://app.example.com/cb",
        "code_challenge": challenge, "code_challenge_method": "S256",
        "state": "xyz",
    }, follow_redirects=False)
    assert r.status_code == 400
    assert r.json()["error"] == "invalid_client"


@pytest.mark.asyncio
async def test_authorize_valid_redirects_to_consent(client):
    client_id = await _register_client(client)
    _, challenge = _pkce_pair()
    r = await client.get("/api/v1/auth/oauth/authorize", params={
        "response_type": "code", "client_id": client_id,
        "redirect_uri": "https://app.example.com/cb",
        "code_challenge": challenge, "code_challenge_method": "S256",
        "state": "xyz",
    }, follow_redirects=False)
    assert r.status_code == 302
    loc = r.headers["location"]
    assert "/oauth/authorize?" in loc
    assert "client_id=" in loc


# ── Full PKCE flow + MCP use ────────────────────────────────────────

async def _full_grant(client, db_session) -> tuple[dict, str]:
    """Run register → decision → token. Returns (token_response, user_id)."""
    user = await _make_employee(db_session)
    await db_session.commit()
    jwt = create_access_token(str(user.id), ["employee"])
    client_id = await _register_client(client)
    verifier, challenge = _pkce_pair()
    redirect_uri = "https://app.example.com/cb"

    decision = await client.post(
        "/api/v1/auth/oauth/authorize/decision",
        json={
            "client_id": client_id, "redirect_uri": redirect_uri,
            "code_challenge": challenge, "code_challenge_method": "S256",
            "scope": "mcp", "state": "st", "approve": True,
        },
        headers={"Authorization": f"Bearer {jwt}"},
    )
    assert decision.status_code == 200, decision.text
    redirect_to = decision.json()["data"]["redirect_to"]
    code = parse_qs(urlparse(redirect_to).query)["code"][0]

    tok = await client.post("/api/v1/auth/oauth/token", data={
        "grant_type": "authorization_code", "code": code,
        "redirect_uri": redirect_uri, "client_id": client_id,
        "code_verifier": verifier,
    })
    assert tok.status_code == 200, tok.text
    return tok.json(), str(user.id)


@pytest.mark.asyncio
async def test_full_flow_token_calls_mcp(client, db_session):
    tokens, user_id = await _full_grant(client, db_session)
    assert tokens["token_type"] == "Bearer"
    assert tokens["scope"] == "mcp"
    access = tokens["access_token"]

    # init then whoami via the real MCP endpoint, authenticated by the
    # OAuth access token (NOT a PAT).
    init = await client.post("/api/v1/mcp/", headers={
        **_MCP_HEADERS, "Authorization": f"Bearer {access}",
    }, json={
        "jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {
            "protocolVersion": "2025-06-18", "capabilities": {},
            "clientInfo": {"name": "test", "version": "0"},
        },
    })
    assert init.status_code == 200, init.text
    payload = _parse_sse(init.text)
    assert payload["result"]["serverInfo"]["name"] == "helix"

    call = await client.post("/api/v1/mcp/", headers={
        **_MCP_HEADERS, "Authorization": f"Bearer {access}",
    }, json={
        "jsonrpc": "2.0", "id": 2, "method": "tools/call",
        "params": {"name": "whoami", "arguments": {}},
    })
    assert call.status_code == 200, call.text
    result = _parse_sse(call.text)
    assert "result" in result, result


@pytest.mark.asyncio
async def test_refresh_rotation_single_use(client, db_session):
    tokens, _ = await _full_grant(client, db_session)
    refresh = tokens["refresh_token"]
    client_id = __import__("jwt").decode(refresh, options={"verify_signature": False})["azp"]

    r1 = await client.post("/api/v1/auth/oauth/token", data={
        "grant_type": "refresh_token", "refresh_token": refresh,
        "client_id": client_id,
    })
    assert r1.status_code == 200, r1.text
    assert r1.json()["access_token"]

    # Replaying the SAME refresh token must fail (rotation = single use).
    r2 = await client.post("/api/v1/auth/oauth/token", data={
        "grant_type": "refresh_token", "refresh_token": refresh,
        "client_id": client_id,
    })
    assert r2.status_code == 400
    assert r2.json()["error"] == "invalid_grant"


@pytest.mark.asyncio
async def test_force_logout_severs_refresh(client, db_session):
    """An admin force-logout must kill the OAuth refresh chain: a refresh
    token issued BEFORE the cutoff cannot rotate into fresh access tokens
    (the "force_logout severs every credential surface" invariant)."""
    from datetime import UTC, datetime, timedelta

    tokens, user_id = await _full_grant(client, db_session)
    refresh = tokens["refresh_token"]
    client_id = __import__("jwt").decode(
        refresh, options={"verify_signature": False},
    )["azp"]

    # Admin force-logs-out the owner AFTER the refresh token was minted.
    person = await db_session.get(IdentityPerson, uuid.UUID(user_id))
    person.force_logout_at = datetime.now(UTC) + timedelta(hours=1)
    await db_session.commit()

    r = await client.post("/api/v1/auth/oauth/token", data={
        "grant_type": "refresh_token", "refresh_token": refresh,
        "client_id": client_id,
    })
    assert r.status_code == 400
    assert r.json()["error"] == "invalid_grant"


@pytest.mark.asyncio
async def test_tampered_verifier_rejected(client, db_session):
    user = await _make_employee(db_session)
    await db_session.commit()
    jwt = create_access_token(str(user.id), ["employee"])
    client_id = await _register_client(client)
    _, challenge = _pkce_pair()
    redirect_uri = "https://app.example.com/cb"

    decision = await client.post(
        "/api/v1/auth/oauth/authorize/decision",
        json={
            "client_id": client_id, "redirect_uri": redirect_uri,
            "code_challenge": challenge, "code_challenge_method": "S256",
            "scope": "mcp", "state": "st", "approve": True,
        },
        headers={"Authorization": f"Bearer {jwt}"},
    )
    code = parse_qs(urlparse(decision.json()["data"]["redirect_to"]).query)["code"][0]

    # Wrong verifier (does not hash to the stored challenge).
    tok = await client.post("/api/v1/auth/oauth/token", data={
        "grant_type": "authorization_code", "code": code,
        "redirect_uri": redirect_uri, "client_id": client_id,
        "code_verifier": secrets.token_urlsafe(64)[:96],
    })
    assert tok.status_code == 400
    assert tok.json()["error"] == "invalid_grant"


async def _make_admin(db_session) -> tuple[IdentityPerson, str]:
    user = await _make_employee(db_session)
    role_id_row = await db_session.execute(
        text("SELECT id FROM auth_roles WHERE name = 'global_admin'")
    )
    role_id = role_id_row.scalar_one()
    if isinstance(role_id, str):
        role_id = uuid.UUID(role_id)
    db_session.add(AuthUserRole(user_id=user.id, role_id=role_id))
    await db_session.flush()
    await db_session.commit()
    return user, create_access_token(str(user.id), ["employee", "global_admin"])


async def _mcp_whoami(client, access: str):
    return await client.post("/api/v1/mcp/", headers={
        **_MCP_HEADERS, "Authorization": f"Bearer {access}",
    }, json={
        "jsonrpc": "2.0", "id": 9, "method": "tools/call",
        "params": {"name": "whoami", "arguments": {}},
    })


@pytest.mark.asyncio
async def test_my_connections_lists_grant(client, db_session):
    tokens, user_id = await _full_grant(client, db_session)
    jwt = create_access_token(user_id, ["employee"])
    r = await client.get("/api/v1/auth/oauth/connections", headers={
        "Authorization": f"Bearer {jwt}"})
    assert r.status_code == 200, r.text
    data = r.json()["data"]
    assert len(data) == 1
    assert data[0]["client_name"] == "Test Connector"


@pytest.mark.asyncio
async def test_self_revoke_blocks_mcp(client, db_session):
    """User self-revoke takes effect on the next MCP request."""
    tokens, user_id = await _full_grant(client, db_session)
    access = tokens["access_token"]
    client_id = __import__("jwt").decode(
        tokens["refresh_token"], options={"verify_signature": False})["azp"]
    jwt = create_access_token(user_id, ["employee"])

    # Works before revoke.
    assert (await _mcp_whoami(client, access)).status_code == 200
    # Revoke own connection.
    rv = await client.post(
        f"/api/v1/auth/oauth/connections/{client_id}/revoke",
        headers={"Authorization": f"Bearer {jwt}"})
    assert rv.status_code == 200 and rv.json()["data"]["revoked"] is True
    # Now the existing access token is rejected immediately.
    assert (await _mcp_whoami(client, access)).status_code == 401
    # And refresh can't rotate either.
    r2 = await client.post("/api/v1/auth/oauth/token", data={
        "grant_type": "refresh_token", "refresh_token": tokens["refresh_token"],
        "client_id": client_id})
    assert r2.status_code == 400


@pytest.mark.asyncio
async def test_admin_disable_blocks_then_enable(client, db_session):
    tokens, _ = await _full_grant(client, db_session)
    access = tokens["access_token"]
    client_id = __import__("jwt").decode(
        tokens["refresh_token"], options={"verify_signature": False})["azp"]
    _admin, admin_jwt = await _make_admin(db_session)

    assert (await _mcp_whoami(client, access)).status_code == 200
    # Admin sees the connection.
    lst = await client.get("/api/v1/auth/oauth/admin/connections", headers={
        "Authorization": f"Bearer {admin_jwt}"})
    assert lst.status_code == 200 and len(lst.json()["data"]) >= 1
    # Disable the connector org-wide → existing token blocked.
    d = await client.post(
        f"/api/v1/auth/oauth/admin/clients/{client_id}/disable",
        headers={"Authorization": f"Bearer {admin_jwt}"})
    assert d.status_code == 200
    assert (await _mcp_whoami(client, access)).status_code == 401
    # Re-enable → works again.
    e = await client.post(
        f"/api/v1/auth/oauth/admin/clients/{client_id}/enable",
        headers={"Authorization": f"Bearer {admin_jwt}"})
    assert e.status_code == 200
    assert (await _mcp_whoami(client, access)).status_code == 200


@pytest.mark.asyncio
async def test_admin_connections_requires_admin(client, db_session):
    user = await _make_employee(db_session)
    await db_session.commit()
    jwt = create_access_token(str(user.id), ["employee"])
    r = await client.get("/api/v1/auth/oauth/admin/connections", headers={
        "Authorization": f"Bearer {jwt}"})
    assert r.status_code == 403


@pytest.mark.asyncio
async def test_wrong_audience_token_rejected_by_mcp(client, db_session):
    """A token whose ``aud`` is a different resource must not unlock MCP."""
    user = await _make_employee(db_session)
    await db_session.commit()
    bad, _ = create_mcp_access_token(
        str(user.id),
        audience="https://someone-else.example.com/api/v1/mcp",
        client_id="helix_mcp_x",
    )
    r = await client.post("/api/v1/mcp/", headers={
        **_MCP_HEADERS, "Authorization": f"Bearer {bad}",
    }, json={"jsonrpc": "2.0", "id": 1, "method": "tools/list", "params": {}})
    assert r.status_code == 401


@pytest.mark.asyncio
async def test_record_grant_is_idempotent_upsert(client, db_session):
    """``record_grant`` must be a race-safe upsert: re-recording the same
    (user, client) never trips ``uq_oauth_grant_user_client`` and instead
    refreshes the row + clears any prior revocation (re-authorize). Guards
    the concurrent first-time-authorize 500 that the SELECT-then-INSERT
    form risked."""
    from app.auth.oauth import service as oauth_service
    from app.auth.oauth.models import OAuthClient

    user = await _make_employee(db_session)
    db_session.add(OAuthClient(
        id=uuid.uuid4(), client_id="helix_mcp_upsert", client_name="Upsert",
        redirect_uris=["https://example.com/cb"],
    ))
    await db_session.commit()

    # First authorization.
    await oauth_service.record_grant(
        db_session, user_id=user.id, client_id="helix_mcp_upsert")
    await db_session.commit()
    # User revokes.
    assert await oauth_service.revoke_grant(
        db_session, user_id=user.id, client_id="helix_mcp_upsert", by="user")
    await db_session.commit()
    assert not await oauth_service.is_grant_active(
        db_session, user_id=user.id, client_id="helix_mcp_upsert")

    # Re-authorize via the SAME path — must not raise and must re-activate.
    await oauth_service.record_grant(
        db_session, user_id=user.id, client_id="helix_mcp_upsert")
    await db_session.commit()
    assert await oauth_service.is_grant_active(
        db_session, user_id=user.id, client_id="helix_mcp_upsert")

    # Exactly one row for the pair (no duplicate insert).
    cnt = (await db_session.execute(text(
        "SELECT COUNT(*) FROM auth_oauth_grants "
        "WHERE user_id = :u AND client_id = 'helix_mcp_upsert'"
    ), {"u": str(user.id)})).scalar()
    assert cnt == 1


@pytest.mark.asyncio
async def test_code_exchange_rejected_after_client_disabled(client, db_session):
    """Admin kill-switch flipped AFTER a code is issued but BEFORE it is
    exchanged must block the exchange (invalid_grant), not mint tokens.

    Defence-in-depth alongside the per-request middleware grant check:
    a code is only issued for an active client, but within the short code
    TTL an admin disabling the client must also fail the in-flight exchange.
    """
    from app.auth.oauth import service as oauth_service

    user = await _make_employee(db_session)
    await db_session.commit()
    jwt = create_access_token(str(user.id), ["employee"])
    client_id = await _register_client(client)
    verifier, challenge = _pkce_pair()
    redirect_uri = "https://app.example.com/cb"

    decision = await client.post(
        "/api/v1/auth/oauth/authorize/decision",
        json={
            "client_id": client_id, "redirect_uri": redirect_uri,
            "code_challenge": challenge, "code_challenge_method": "S256",
            "scope": "mcp", "state": "st", "approve": True,
        },
        headers={"Authorization": f"Bearer {jwt}"},
    )
    assert decision.status_code == 200, decision.text
    code = parse_qs(urlparse(decision.json()["data"]["redirect_to"]).query)["code"][0]

    # Admin kill-switch AFTER the code was issued.
    assert await oauth_service.set_client_disabled(
        db_session, client_id=client_id, disabled=True)
    await db_session.commit()

    tok = await client.post("/api/v1/auth/oauth/token", data={
        "grant_type": "authorization_code", "code": code,
        "redirect_uri": redirect_uri, "client_id": client_id,
        "code_verifier": verifier,
    })
    assert tok.status_code == 400, tok.text
    assert tok.json().get("error") == "invalid_grant"