"""MCP OAuth 2.1 endpoints — discovery, DCR, authorize, token, revoke.

These are *standards-compliant OAuth surfaces*, so unlike the rest of the
Helix API they do NOT use the ``{success, data, meta, errors}`` envelope:
discovery returns RFC 8414 / 9728 metadata, the token endpoint returns
the RFC 6749 token response, and errors return ``{error,
error_description}``. The one exception is the consent *decision*
endpoint, which is called by our own SPA and uses the Helix envelope.

Discovery endpoints are mounted at the domain root (``/.well-known/...``)
via ``wellknown_router`` (no ``/api/v1`` prefix). nginx proxies
``/.well-known/oauth`` to the backend (see ``web/nginx.conf``).
"""
from __future__ import annotations

import logging
from urllib.parse import urlencode

from fastapi import APIRouter, Depends, Request
from fastapi.responses import JSONResponse, RedirectResponse
from sqlalchemy.ext.asyncio import AsyncSession

from app.auth.oauth import service as oauth_service
from app.auth.oauth.metadata import (
    authorization_server_metadata,
    mcp_resource_uri,
    origin_from_request,
    protected_resource_metadata,
)
from app.auth.oauth.service import OAuthError
from app.auth.service import is_effectively_departed
from app.core.database import get_db
from app.core.dependencies import (
    _load_user_with_roles,
    check_login_rate_limit,
    get_current_user,
    get_real_ip,
    require_admin,
)
from app.core.distributed_lock import LockBusy, distributed_lock
from app.core.response import success_response
from app.core.security import (
    create_mcp_access_token,
    create_mcp_refresh_token,
    decode_mcp_refresh_token,
)

logger = logging.getLogger(__name__)

wellknown_router = APIRouter(tags=["oauth"])
router = APIRouter(prefix="/api/v1/auth/oauth", tags=["oauth"])

# Cache-friendly, no-credential metadata. Allow any origin to read it —
# discovery documents are public by design.
_METADATA_HEADERS = {
    "Cache-Control": "public, max-age=3600",
    "Access-Control-Allow-Origin": "*",
}


def _oauth_error(err: OAuthError) -> JSONResponse:
    return JSONResponse(
        status_code=err.status,
        content={"error": err.error, "error_description": err.description},
    )


# ── Discovery (RFC 9728 / RFC 8414) ─────────────────────────────────

@wellknown_router.get("/.well-known/oauth-protected-resource")
@wellknown_router.get("/.well-known/oauth-protected-resource/api/v1/mcp")
async def protected_resource(request: Request) -> JSONResponse:
    return JSONResponse(
        content=protected_resource_metadata(origin_from_request(request)),
        headers=_METADATA_HEADERS,
    )


@wellknown_router.get("/.well-known/oauth-authorization-server")
@wellknown_router.get("/.well-known/oauth-authorization-server/api/v1/mcp")
async def authorization_server(request: Request) -> JSONResponse:
    return JSONResponse(
        content=authorization_server_metadata(origin_from_request(request)),
        headers=_METADATA_HEADERS,
    )


# ── Dynamic client registration (RFC 7591) ──────────────────────────

@router.post("/register")
async def register(
    request: Request,
    db: AsyncSession = Depends(get_db),
    _rate: None = Depends(check_login_rate_limit),
) -> JSONResponse:
    try:
        body = await request.json()
    except Exception:
        return _oauth_error(
            OAuthError("invalid_client_metadata", "request body must be JSON")
        )
    if not isinstance(body, dict):
        return _oauth_error(
            OAuthError("invalid_client_metadata", "request body must be a JSON object")
        )
    source_ip = get_real_ip(request)
    try:
        # Serialize cleanup + quota check + insert across API workers. The
        # commit stays inside the lock so the next registrant observes it.
        async with distributed_lock("oauth_dynamic_client_registration", ttl=30):
            await oauth_service.prune_stale_ungranted_clients(db)
            await oauth_service.enforce_registration_quota(
                db, created_ip=source_ip,
            )
            client = await oauth_service.register_client(
                db,
                client_name=body.get("client_name"),
                redirect_uris=body.get("redirect_uris") or [],
                grant_types=body.get("grant_types"),
                response_types=body.get("response_types"),
                scope=body.get("scope"),
                token_endpoint_auth_method=body.get("token_endpoint_auth_method"),
                created_ip=source_ip,
                created_user_agent=request.headers.get("User-Agent"),
            )
            await db.commit()
    except LockBusy:
        await db.rollback()
        return _oauth_error(OAuthError(
            "temporarily_unavailable",
            "another dynamic registration is in progress; retry shortly",
            status=429,
        ))
    except OAuthError as err:
        await db.rollback()
        return _oauth_error(err)
    return JSONResponse(
        status_code=201,
        content=oauth_service.client_registration_payload(client),
    )


# ── Authorization endpoint ──────────────────────────────────────────

def _redirect_with(redirect_uri: str, params: dict) -> RedirectResponse:
    sep = "&" if "?" in redirect_uri else "?"
    return RedirectResponse(
        url=f"{redirect_uri}{sep}{urlencode(params)}", status_code=302,
    )


@router.get("/authorize", response_model=None)
async def authorize(
    request: Request,
    db: AsyncSession = Depends(get_db),
) -> JSONResponse | RedirectResponse:
    """Validate the authorization request, then bounce the browser to the
    SPA consent page. The user authenticates there (existing SSO) and the
    SPA POSTs the consent decision back to ``/authorize/decision``.

    Per OAuth 2.1 open-redirect rules we validate ``client_id`` +
    ``redirect_uri`` HERE and refuse with a 400 (never a redirect) if
    either is unknown — only an exact-registered redirect_uri is ever
    used as a redirect target.
    """
    q = request.query_params
    client_id = q.get("client_id", "")
    redirect_uri = q.get("redirect_uri", "")

    client = await oauth_service.get_active_client(db, client_id)
    if client is None:
        return _oauth_error(OAuthError("invalid_client", "unknown client_id"))
    if redirect_uri not in (client.redirect_uris or []):
        return _oauth_error(
            OAuthError("invalid_request", "redirect_uri not registered for this client")
        )

    state = q.get("state", "")
    # Beyond this point redirect_uri is validated → errors go back to it.
    if q.get("response_type") != "code":
        return _redirect_with(redirect_uri, {
            "error": "unsupported_response_type", "state": state,
        })
    code_challenge = q.get("code_challenge", "")
    method = q.get("code_challenge_method", "S256")
    if not code_challenge or method != "S256":
        return _redirect_with(redirect_uri, {
            "error": "invalid_request",
            "error_description": "PKCE S256 code_challenge required",
            "state": state,
        })

    # Bounce to the SPA consent route, preserving the (validated) request.
    origin = origin_from_request(request)
    consent_params = {
        "client_id": client_id,
        "redirect_uri": redirect_uri,
        "code_challenge": code_challenge,
        "code_challenge_method": method,
        "scope": q.get("scope", "mcp"),
        "state": state,
        "resource": q.get("resource", mcp_resource_uri(origin)),
    }
    return RedirectResponse(
        url=f"{origin}/oauth/authorize?{urlencode(consent_params)}",
        status_code=302,
    )


@router.post("/authorize/decision")
async def authorize_decision(
    request: Request,
    db: AsyncSession = Depends(get_db),
    current_user=Depends(get_current_user),
):
    """Consent decision from the SPA (the user is logged in). Returns the
    ``redirect_to`` URL the SPA should navigate the browser to."""
    body = await request.json()
    client_id = body.get("client_id", "")
    redirect_uri = body.get("redirect_uri", "")
    state = body.get("state", "")
    approve = bool(body.get("approve"))

    client = await oauth_service.get_active_client(db, client_id)
    if client is None or redirect_uri not in (client.redirect_uris or []):
        # Don't bounce to an unvalidated URI — surface to the SPA instead.
        return success_response(
            data={"error": "invalid_request"}, request=request,
        )

    if not approve:
        return success_response(
            data={"redirect_to": f"{redirect_uri}?" + urlencode({
                "error": "access_denied", "state": state,
            })},
            request=request,
        )

    # Re-validate PKCE at the point the code is actually MINTED — the GET
    # /authorize check (lines above) gates the consent redirect, but this
    # POST is a separate request that could be called directly. Without
    # this, a decision POST with no ``code_challenge`` mints a code stored
    # with an empty challenge; ``verify_pkce`` then rejects every redeem,
    # yielding a confusing dead-end auth failure instead of a clear error.
    code_challenge = body.get("code_challenge", "")
    code_challenge_method = body.get("code_challenge_method", "S256")
    if not code_challenge or code_challenge_method != "S256":
        return success_response(
            data={"redirect_to": f"{redirect_uri}?" + urlencode({
                "error": "invalid_request",
                "error_description": "PKCE S256 code_challenge required",
                "state": state,
            })},
            request=request,
        )

    origin = origin_from_request(request)
    resource = body.get("resource") or mcp_resource_uri(origin)
    code = await oauth_service.issue_authorization_code(
        client_id=client_id,
        user_id=str(current_user.id),
        redirect_uri=redirect_uri,
        code_challenge=code_challenge,
        code_challenge_method=code_challenge_method,
        scope=body.get("scope", "mcp"),
        resource=resource,
    )
    await _audit(
        request, db, actor_user_id=str(current_user.id),
        action="mcp_oauth.authorization_granted",
        detail={"client_id": client_id, "client_name": client.client_name},
    )
    return success_response(
        data={"redirect_to": f"{redirect_uri}?" + urlencode({
            "code": code, "state": state,
        })},
        request=request,
    )


# ── Token endpoint ──────────────────────────────────────────────────

async def _read_token_params(request: Request) -> dict:
    """Token requests are form-encoded per OAuth, but accept JSON too."""
    ctype = request.headers.get("content-type", "")
    if "application/json" in ctype:
        try:
            data = await request.json()
            return data if isinstance(data, dict) else {}
        except Exception:
            return {}
    form = await request.form()
    return {k: v for k, v in form.items()}


@router.post("/token")
async def token(
    request: Request,
    db: AsyncSession = Depends(get_db),
) -> JSONResponse:
    params = await _read_token_params(request)
    grant_type = params.get("grant_type")
    try:
        if grant_type == "authorization_code":
            return await _grant_authorization_code(request, db, params)
        if grant_type == "refresh_token":
            return await _grant_refresh_token(request, db, params)
        raise OAuthError("unsupported_grant_type", f"unsupported grant_type: {grant_type}")
    except OAuthError as err:
        return _oauth_error(err)


async def _active_user_or_error(db: AsyncSession, user_id: str):
    from uuid import UUID
    try:
        user = await _load_user_with_roles(db, UUID(str(user_id)))
    except (ValueError, TypeError):
        user = None
    if user is None or is_effectively_departed(user):
        raise OAuthError("invalid_grant", "token owner is inactive", status=400)
    return user


async def _grant_authorization_code(
    request: Request, db: AsyncSession, params: dict,
) -> JSONResponse:
    code = params.get("code", "")
    client_id = params.get("client_id", "")
    redirect_uri = params.get("redirect_uri", "")
    code_verifier = params.get("code_verifier", "")

    stored = await oauth_service.consume_authorization_code(code)
    if stored is None:
        raise OAuthError("invalid_grant", "authorization code is invalid or expired")
    if stored.client_id != client_id:
        raise OAuthError("invalid_grant", "client_id mismatch")
    if stored.redirect_uri != redirect_uri:
        raise OAuthError("invalid_grant", "redirect_uri mismatch")
    if not oauth_service.verify_pkce(
        code_verifier, stored.code_challenge, stored.code_challenge_method,
    ):
        raise OAuthError("invalid_grant", "PKCE verification failed")

    # Re-confirm the client is still enabled at exchange time. The code is
    # only ISSUED for an active client (authorize gates on get_active_client),
    # but within the short code TTL an admin kill-switch (disabled_at) must
    # also block an in-flight exchange — defence in depth alongside the
    # per-request middleware grant check that already rejects the minted token.
    if await oauth_service.get_active_client(db, client_id) is None:
        raise OAuthError("invalid_grant", "client is no longer active")

    user = await _active_user_or_error(db, stored.user_id)
    # Establish the durable grant (this is the moment a connection truly
    # forms — recorded here, not at consent, so an abandoned consent
    # leaves no phantom "connected app").
    await oauth_service.record_grant(db, user_id=user.id, client_id=client_id)
    await db.commit()
    audience = stored.resource or mcp_resource_uri(origin_from_request(request))
    resp = await _issue_token_pair(
        user_id=str(user.id), client_id=client_id, audience=audience,
    )
    await _audit(
        request, db, actor_user_id=str(user.id),
        action="mcp_oauth.token_issued",
        detail={"client_id": client_id, "grant": "authorization_code"},
    )
    return resp


async def _grant_refresh_token(
    request: Request, db: AsyncSession, params: dict,
) -> JSONResponse:
    client_id = params.get("client_id", "")
    refresh_token = params.get("refresh_token", "")
    try:
        payload = decode_mcp_refresh_token(refresh_token)
    except Exception:
        raise OAuthError("invalid_grant", "refresh token is invalid or expired")
    jti = payload.get("jti", "")
    # Single-use consume (rotation): a replayed refresh token finds no jti.
    stored = await oauth_service.consume_refresh_jti(jti)
    if stored is None:
        raise OAuthError("invalid_grant", "refresh token already used or expired")
    if payload.get("azp") != client_id or stored.get("client_id") != client_id:
        raise OAuthError("invalid_grant", "client_id mismatch")

    user = await _active_user_or_error(db, stored.get("user_id"))
    # Force-logout parity (the "force_logout severs every credential
    # surface" invariant): a refresh token issued BEFORE the owner's
    # force_logout cutoff is dead. Without this, an admin force-logout
    # (suspected compromise) is outlived by a pre-cutoff refresh token
    # that keeps minting fresh access tokens — each with a post-cutoff
    # iat that the MCP middleware's access-token force_logout check then
    # waves through. Refuse the rotation so the chain breaks within one
    # access-token lifetime.
    from datetime import UTC, datetime

    from app.core.dependencies import _enforce_force_logout
    from app.core.exceptions import UnauthorizedError
    try:
        refresh_iat = datetime.fromtimestamp(int(payload["iat"]), tz=UTC)
        _enforce_force_logout(user, refresh_iat)
    except UnauthorizedError:
        raise OAuthError("invalid_grant", "session was terminated", status=400)
    except (KeyError, ValueError, TypeError):
        pass
    # Grant must still be live (not user-revoked, client not admin-disabled).
    if not await oauth_service.is_grant_active(db, user_id=user.id, client_id=client_id):
        raise OAuthError("invalid_grant", "authorization was revoked", status=400)
    await oauth_service.touch_grant(db, user_id=user.id, client_id=client_id)
    await db.commit()
    audience = mcp_resource_uri(origin_from_request(request))
    resp = await _issue_token_pair(
        user_id=str(user.id), client_id=client_id, audience=audience,
    )
    await _audit(
        request, db, actor_user_id=str(user.id),
        action="mcp_oauth.token_issued",
        detail={"client_id": client_id, "grant": "refresh_token"},
    )
    return resp


async def _issue_token_pair(
    *, user_id: str, client_id: str, audience: str,
) -> JSONResponse:
    access_token, expires_in = create_mcp_access_token(
        user_id, audience=audience, client_id=client_id,
    )
    refresh_token, jti = create_mcp_refresh_token(user_id, client_id=client_id)
    await oauth_service.store_refresh_jti(jti, user_id=user_id, client_id=client_id)
    return JSONResponse(
        status_code=200,
        headers={"Cache-Control": "no-store", "Pragma": "no-cache"},
        content={
            "access_token": access_token,
            "token_type": "Bearer",
            "expires_in": expires_in,
            "refresh_token": refresh_token,
            "scope": "mcp",
        },
    )


# ── Revocation (RFC 7009) ───────────────────────────────────────────

@router.post("/revoke")
async def revoke(request: Request) -> JSONResponse:
    params = await _read_token_params(request)
    tok = params.get("token", "")
    # Best-effort: revoke the refresh jti. Per RFC 7009 the endpoint
    # returns 200 even for an unknown / already-invalid token.
    try:
        payload = decode_mcp_refresh_token(tok)
        await oauth_service.revoke_refresh_jti(payload.get("jti", ""))
    except Exception:
        pass
    return JSONResponse(status_code=200, content={})


# ── Governance: self-service "my connections" ──────────────────────

@router.get("/connections")
async def my_connections(
    request: Request,
    db: AsyncSession = Depends(get_db),
    current_user=Depends(get_current_user),
):
    """Active MCP connections the caller has authorized (self-scoped)."""
    data = await oauth_service.list_user_grants(db, user_id=current_user.id)
    return success_response(data=data, request=request)


@router.post("/connections/{client_id}/revoke")
async def revoke_my_connection(
    client_id: str,
    request: Request,
    db: AsyncSession = Depends(get_db),
    current_user=Depends(get_current_user),
):
    """Revoke the caller's own grant for a connector. Takes effect on the
    next MCP request (middleware re-checks the grant per call)."""
    revoked = await oauth_service.revoke_grant(
        db, user_id=current_user.id, client_id=client_id, by="user",
    )
    await db.commit()
    if revoked:
        await _audit(
            request, db, actor_user_id=str(current_user.id),
            action="mcp_oauth.connection_revoked",
            detail={"client_id": client_id, "by": "user"},
        )
    return success_response(data={"revoked": revoked}, request=request)


# ── Governance: admin "connected apps" ─────────────────────────────

@router.get("/admin/connections")
async def admin_connections(
    request: Request,
    db: AsyncSession = Depends(get_db),
    _admin=Depends(require_admin()),
):
    """Every MCP OAuth connection across the org (who authorized which
    connector, last use, revocation / disable state). Admin-only."""
    data = await oauth_service.list_all_grants(db)
    return success_response(data=data, request=request)


@router.post("/admin/clients/{client_id}/disable")
async def admin_disable_client(
    client_id: str,
    request: Request,
    db: AsyncSession = Depends(get_db),
    admin=Depends(require_admin()),
):
    """Kill-switch a connector org-wide: every existing grant stops on the
    next MCP request, and no new authorization can complete."""
    ok = await oauth_service.set_client_disabled(db, client_id=client_id, disabled=True)
    await db.commit()
    if ok:
        await _audit(
            request, db, actor_user_id=str(admin.id),
            action="mcp_oauth.client_disabled",
            detail={"client_id": client_id},
        )
    return success_response(data={"ok": ok}, request=request)


@router.post("/admin/clients/{client_id}/enable")
async def admin_enable_client(
    client_id: str,
    request: Request,
    db: AsyncSession = Depends(get_db),
    admin=Depends(require_admin()),
):
    """Re-enable a previously disabled connector (grants resume)."""
    ok = await oauth_service.set_client_disabled(db, client_id=client_id, disabled=False)
    await db.commit()
    if ok:
        await _audit(
            request, db, actor_user_id=str(admin.id),
            action="mcp_oauth.client_enabled",
            detail={"client_id": client_id},
        )
    return success_response(data={"ok": ok}, request=request)


async def _audit(
    request: Request, db: AsyncSession, *, actor_user_id: str, action: str,
    detail: dict,
) -> None:
    """Best-effort audit row for an auth-server event (own session)."""
    try:
        from uuid import uuid4

        from app.audit.service import AuditService
        from app.core import database as _db

        async with _db.async_session_factory() as audit_db:
            await AuditService(audit_db).log(
                trace_id=str(uuid4()),
                actor_type="user",
                actor_user_id=actor_user_id,
                action=action,
                status="success",
                resource_type="mcp_oauth",
                resource_id=detail.get("client_id", "oauth"),
                detail=detail,
                ip_address=get_real_ip(request),
                user_agent=request.headers.get("User-Agent"),
            )
            await audit_db.commit()
    except Exception:  # noqa: BLE001
        logger.exception("mcp_oauth audit write failed (%s)", action)


__all__ = ["router", "wellknown_router"]