"""MCP OAuth 2.1 Authorization Server — core logic.

Responsibilities:
  * RFC 7591 dynamic client registration (public clients only).
  * Authorization-code issue / single-use consume (Redis, short TTL).
  * PKCE (RFC 7636, S256) challenge ↔ verifier validation.
  * Refresh-token rotation (single-use jti in Redis) + revocation.

Access/refresh JWTs themselves are minted in :mod:`app.core.security`
(``create_mcp_access_token`` / ``create_mcp_refresh_token``); this module
owns the *state* (registered clients, live codes, refresh jtis) and the
validation rules around them.
"""
from __future__ import annotations

import base64
import hashlib
import json
import logging
import re
import secrets
import uuid
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from urllib.parse import urlparse

from sqlalchemy import delete, func, select
from sqlalchemy.dialects.mysql import insert as mysql_insert
from sqlalchemy.ext.asyncio import AsyncSession

from app.auth.oauth.models import OAuthClient, OAuthGrant
from app.core.config import settings
from app.core.redis_client import get_redis

logger = logging.getLogger(__name__)

_CODE_PREFIX = "mcp_oauth:code:"
_REFRESH_PREFIX = "mcp_oauth:refresh:"
_CLIENT_ID_PREFIX = "helix_mcp_"

# Dynamic registration is intentionally unauthenticated for MCP client
# interoperability. Keep that public surface small and bounded: registrations
# that never produce a grant are disposable bootstrap state, not a durable
# tenant directory.
MAX_REDIRECT_URIS = 10
MAX_REDIRECT_URI_LENGTH = 2048
MAX_UNGRANTED_CLIENTS = 1000
MAX_UNGRANTED_CLIENTS_PER_IP = 50
UNGRANTED_CLIENT_TTL_DAYS = 7
_SUPPORTED_GRANT_TYPES = frozenset({"authorization_code", "refresh_token"})
_SUPPORTED_RESPONSE_TYPES = frozenset({"code"})

# PKCE verifier: 43–128 chars from the unreserved set (RFC 7636 §4.1).
_VERIFIER_RE = re.compile(r"^[A-Za-z0-9._~-]{43,128}$")


class OAuthError(Exception):
    """OAuth 2.0 protocol error → rendered as ``{error, error_description}``.

    ``status`` is the HTTP status the token/register endpoints return.
    """

    def __init__(self, error: str, description: str, *, status: int = 400) -> None:
        super().__init__(description)
        self.error = error
        self.description = description
        self.status = status


# ── PKCE ────────────────────────────────────────────────────────────

def verify_pkce(code_verifier: str, code_challenge: str, method: str) -> bool:
    """Return True iff ``code_verifier`` proves possession of the secret
    behind ``code_challenge``. Only S256 is supported (plain is forbidden
    by the MCP spec / OAuth 2.1 for public clients)."""
    if not code_verifier or not _VERIFIER_RE.match(code_verifier):
        return False
    if (method or "S256") != "S256":
        return False
    digest = hashlib.sha256(code_verifier.encode("ascii")).digest()
    expected = base64.urlsafe_b64encode(digest).decode("ascii").rstrip("=")
    return secrets.compare_digest(expected, code_challenge or "")


# ── Redirect-URI validation ─────────────────────────────────────────

def _is_valid_redirect_uri(uri: str) -> bool:
    """OAuth 2.1: redirect URIs must be absolute and either HTTPS or a
    loopback ``http://localhost`` / ``127.0.0.1`` (native-app loopback)."""
    try:
        p = urlparse(uri)
    except ValueError:
        return False
    if not p.scheme or not p.netloc or p.fragment:
        return False
    if p.scheme == "https":
        return True
    if p.scheme == "http":
        host = (p.hostname or "").lower()
        return host in ("localhost", "127.0.0.1", "::1")
    # Custom app schemes (e.g. ``cursor://``) — allow non-fragment absolute
    # URIs so native MCP clients that use a private-use scheme can register.
    return bool(p.scheme and p.netloc or p.path)


def _validate_string_list(
    value: list[str] | None,
    *,
    field: str,
    default: list[str],
    allowed: frozenset[str],
) -> list[str]:
    items = default if value is None else value
    if not isinstance(items, list) or not items:
        raise OAuthError("invalid_client_metadata", f"{field} must be a non-empty array")
    if any(not isinstance(item, str) for item in items):
        raise OAuthError("invalid_client_metadata", f"{field} entries must be strings")
    normalized = list(dict.fromkeys(item.strip() for item in items))
    if any(not item or item not in allowed for item in normalized):
        raise OAuthError(
            "invalid_client_metadata",
            f"{field} contains an unsupported value",
        )
    return normalized


def _ungranted_client_clause():
    return ~select(OAuthGrant.id).where(
        OAuthGrant.client_id == OAuthClient.client_id,
    ).exists()


async def prune_stale_ungranted_clients(
    db: AsyncSession, *, now: datetime | None = None,
) -> int:
    """Delete expired registrations that never reached user consent."""
    cutoff = (now or datetime.now(UTC)) - timedelta(days=UNGRANTED_CLIENT_TTL_DAYS)
    result = await db.execute(
        delete(OAuthClient).where(
            OAuthClient.created_at < cutoff,
            _ungranted_client_clause(),
        )
    )
    # AsyncSession.execute is typed as generic Result even though DML returns
    # a CursorResult with rowcount at runtime.
    return int(getattr(result, "rowcount", 0) or 0)


async def enforce_registration_quota(
    db: AsyncSession, *, created_ip: str | None,
) -> None:
    """Bound durable anonymous DCR state.

    The router calls this under one global distributed lock and commits before
    releasing it, so concurrent processes cannot both observe quota headroom.
    """
    ungranted = _ungranted_client_clause()
    total = await db.scalar(
        select(func.count()).select_from(OAuthClient).where(ungranted)
    )
    if int(total or 0) >= MAX_UNGRANTED_CLIENTS:
        raise OAuthError(
            "temporarily_unavailable",
            "dynamic client registration capacity is temporarily exhausted",
            status=503,
        )
    if created_ip:
        per_ip = await db.scalar(
            select(func.count()).select_from(OAuthClient).where(
                ungranted,
                OAuthClient.created_ip == created_ip,
            )
        )
        if int(per_ip or 0) >= MAX_UNGRANTED_CLIENTS_PER_IP:
            raise OAuthError(
                "registration_access_denied",
                "too many ungranted clients from this source",
                status=429,
            )


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

async def register_client(
    db: AsyncSession,
    *,
    client_name: str | None,
    redirect_uris: list[str],
    grant_types: list[str] | None,
    response_types: list[str] | None,
    scope: str | None,
    token_endpoint_auth_method: str | None,
    created_ip: str | None = None,
    created_user_agent: str | None = None,
) -> OAuthClient:
    """RFC 7591 dynamic client registration. Public clients only."""
    if client_name is not None and not isinstance(client_name, str):
        raise OAuthError("invalid_client_metadata", "client_name must be a string")
    if not isinstance(redirect_uris, list):
        raise OAuthError("invalid_redirect_uri", "redirect_uris must be an array")
    if len(redirect_uris) > MAX_REDIRECT_URIS:
        raise OAuthError(
            "invalid_redirect_uri",
            f"redirect_uris may contain at most {MAX_REDIRECT_URIS} entries",
        )
    if any(not isinstance(u, str) for u in redirect_uris):
        raise OAuthError("invalid_redirect_uri", "redirect_uris entries must be strings")
    uris = list(dict.fromkeys(u.strip() for u in redirect_uris if u.strip()))
    if not uris:
        raise OAuthError(
            "invalid_redirect_uri", "redirect_uris is required and must be non-empty",
        )
    oversized = [u for u in uris if len(u) > MAX_REDIRECT_URI_LENGTH]
    if oversized:
        raise OAuthError(
            "invalid_redirect_uri",
            f"redirect_uri exceeds {MAX_REDIRECT_URI_LENGTH} characters",
        )
    bad = [u for u in uris if not _is_valid_redirect_uri(u)]
    if bad:
        raise OAuthError(
            "invalid_redirect_uri",
            f"redirect_uri must be https or http://localhost: {bad[0]}",
        )
    if token_endpoint_auth_method is not None and not isinstance(
        token_endpoint_auth_method, str,
    ):
        raise OAuthError(
            "invalid_client_metadata",
            "token_endpoint_auth_method must be a string",
        )
    method = (token_endpoint_auth_method or "none").strip()
    if method != "none":
        # We only issue public clients (PKCE). Reject confidential-client
        # registration rather than silently downgrading it.
        raise OAuthError(
            "invalid_client_metadata",
            "only public clients (token_endpoint_auth_method=none) are supported",
        )
    grants = _validate_string_list(
        grant_types,
        field="grant_types",
        default=["authorization_code", "refresh_token"],
        allowed=_SUPPORTED_GRANT_TYPES,
    )
    if "authorization_code" not in grants:
        raise OAuthError(
            "invalid_client_metadata",
            "grant_types must include authorization_code",
        )
    responses = _validate_string_list(
        response_types,
        field="response_types",
        default=["code"],
        allowed=_SUPPORTED_RESPONSE_TYPES,
    )
    if scope is not None and not isinstance(scope, str):
        raise OAuthError("invalid_client_metadata", "scope must be a string")
    normalized_scope = (scope or "mcp").strip()
    if normalized_scope != "mcp":
        raise OAuthError("invalid_client_metadata", "only scope=mcp is supported")
    client = OAuthClient(
        id=uuid.uuid4(),
        client_id=f"{_CLIENT_ID_PREFIX}{secrets.token_urlsafe(24)}",
        client_name=(client_name or "").strip()[:255] or None,
        redirect_uris=uris,
        grant_types=grants,
        response_types=responses,
        scope=normalized_scope,
        token_endpoint_auth_method="none",
        created_ip=created_ip,
        created_user_agent=(created_user_agent or "")[:512] or None,
    )
    db.add(client)
    await db.flush()
    return client


async def get_active_client(db: AsyncSession, client_id: str) -> OAuthClient | None:
    if not client_id:
        return None
    row = await db.execute(
        select(OAuthClient).where(OAuthClient.client_id == client_id)
    )
    client = row.scalar_one_or_none()
    if client is None or client.disabled_at is not None:
        return None
    return client


def client_registration_payload(client: OAuthClient) -> dict:
    """RFC 7591 registration response body."""
    return {
        "client_id": client.client_id,
        "client_id_issued_at": int(client.created_at.timestamp())
        if client.created_at else None,
        "client_name": client.client_name,
        "redirect_uris": client.redirect_uris,
        "grant_types": client.grant_types,
        "response_types": client.response_types,
        "scope": client.scope,
        "token_endpoint_auth_method": client.token_endpoint_auth_method,
    }


# ── Authorization codes ─────────────────────────────────────────────

@dataclass(slots=True)
class AuthorizationCode:
    client_id: str
    user_id: str
    redirect_uri: str
    code_challenge: str
    code_challenge_method: str
    scope: str
    resource: str


async def issue_authorization_code(
    *,
    client_id: str,
    user_id: str,
    redirect_uri: str,
    code_challenge: str,
    code_challenge_method: str,
    scope: str,
    resource: str,
) -> str:
    """Mint + store a single-use authorization code (Redis, short TTL)."""
    code = secrets.token_urlsafe(32)
    payload = json.dumps({
        "client_id": client_id,
        "user_id": user_id,
        "redirect_uri": redirect_uri,
        "code_challenge": code_challenge,
        "code_challenge_method": code_challenge_method or "S256",
        "scope": scope,
        "resource": resource,
    })
    redis = get_redis()
    await redis.setex(
        f"{_CODE_PREFIX}{code}", settings.MCP_OAUTH_CODE_TTL_SECONDS, payload,
    )
    return code


async def consume_authorization_code(code: str) -> AuthorizationCode | None:
    """Atomic single-use consume (GETDEL) of an authorization code."""
    if not code:
        return None
    redis = get_redis()
    raw = await redis.getdel(f"{_CODE_PREFIX}{code}")
    if raw is None:
        return None
    if isinstance(raw, bytes):
        raw = raw.decode("utf-8")
    try:
        data = json.loads(raw)
    except (ValueError, TypeError):
        return None
    return AuthorizationCode(
        client_id=data["client_id"],
        user_id=data["user_id"],
        redirect_uri=data["redirect_uri"],
        code_challenge=data["code_challenge"],
        code_challenge_method=data.get("code_challenge_method", "S256"),
        scope=data.get("scope", "mcp"),
        resource=data.get("resource", ""),
    )


# ── Refresh-token rotation ──────────────────────────────────────────

async def store_refresh_jti(jti: str, *, user_id: str, client_id: str) -> None:
    redis = get_redis()
    ttl = settings.MCP_OAUTH_REFRESH_TTL_DAYS * 86400
    await redis.setex(
        f"{_REFRESH_PREFIX}{jti}", ttl,
        json.dumps({"user_id": user_id, "client_id": client_id}),
    )


async def consume_refresh_jti(jti: str) -> dict | None:
    """Single-use consume of a refresh jti (rotation). Returns the stored
    ``{user_id, client_id}`` or None if unknown / already used / expired."""
    if not jti:
        return None
    redis = get_redis()
    raw = await redis.getdel(f"{_REFRESH_PREFIX}{jti}")
    if raw is None:
        return None
    if isinstance(raw, bytes):
        raw = raw.decode("utf-8")
    try:
        return json.loads(raw)
    except (ValueError, TypeError):
        return None


async def revoke_refresh_jti(jti: str) -> None:
    if not jti:
        return
    redis = get_redis()
    await redis.delete(f"{_REFRESH_PREFIX}{jti}")


# ── Grants (durable per-user authorization records) ─────────────────

async def record_grant(db: AsyncSession, *, user_id, client_id: str) -> None:
    """Upsert the (user, client) authorization on consent. Re-authorizing
    refreshes ``granted_at`` and clears any prior revocation.

    Atomic INSERT … ON DUPLICATE KEY UPDATE on ``uq_oauth_grant_user_client``
    so two concurrent first-time authorization-code exchanges for the same
    (user, client) can't both INSERT and trip the unique constraint (which
    would surface as an unhandled 500 to the connector). The re-authorize
    branch is idempotent, so collapsing both paths into one upsert is safe.
    """
    now = datetime.now(UTC)
    stmt = mysql_insert(OAuthGrant).values(
        id=str(uuid.uuid4()),
        user_id=user_id, client_id=client_id,
        granted_at=now, last_used_at=now,
    ).on_duplicate_key_update(
        granted_at=now, last_used_at=now,
        revoked_at=None, revoked_by=None,
    )
    await db.execute(stmt)
    await db.flush()


async def touch_grant(db: AsyncSession, *, user_id, client_id: str) -> None:
    """Best-effort bump of ``last_used_at`` (called on each token issue)."""
    row = (await db.execute(
        select(OAuthGrant).where(
            OAuthGrant.user_id == user_id, OAuthGrant.client_id == client_id,
        )
    )).scalar_one_or_none()
    if row is not None:
        row.last_used_at = datetime.now(UTC)
        await db.flush()


async def is_grant_active(db: AsyncSession, *, user_id, client_id: str) -> bool:
    """True iff the user has a live grant for this client AND the client
    is not disabled. Consulted by the MCP middleware + refresh grant so a
    revoked grant / disabled client is immediate logout."""
    if not client_id:
        return False
    row = (await db.execute(
        select(OAuthGrant.revoked_at, OAuthClient.disabled_at)
        .join(OAuthClient, OAuthClient.client_id == OAuthGrant.client_id)
        .where(
            OAuthGrant.user_id == user_id,
            OAuthGrant.client_id == client_id,
        )
    )).first()
    if row is None:
        return False
    grant_revoked, client_disabled = row
    return grant_revoked is None and client_disabled is None


async def revoke_grant(
    db: AsyncSession, *, user_id, client_id: str, by: str,
) -> bool:
    """Revoke a single user's grant for a client. Returns True if a live
    grant was found and revoked."""
    row = (await db.execute(
        select(OAuthGrant).where(
            OAuthGrant.user_id == user_id,
            OAuthGrant.client_id == client_id,
            OAuthGrant.revoked_at.is_(None),
        )
    )).scalar_one_or_none()
    if row is None:
        return False
    row.revoked_at = datetime.now(UTC)
    row.revoked_by = by
    await db.flush()
    return True


async def list_user_grants(db: AsyncSession, *, user_id) -> list[dict]:
    """Active connections for one user (for the self-service list)."""
    rows = (await db.execute(
        select(
            OAuthGrant.client_id, OAuthGrant.granted_at,
            OAuthGrant.last_used_at, OAuthClient.client_name,
            OAuthClient.disabled_at,
        )
        .join(OAuthClient, OAuthClient.client_id == OAuthGrant.client_id)
        .where(OAuthGrant.user_id == user_id, OAuthGrant.revoked_at.is_(None))
        .order_by(OAuthGrant.granted_at.desc())
    )).all()
    return [
        {
            "client_id": r.client_id,
            "client_name": r.client_name,
            "granted_at": r.granted_at,
            "last_used_at": r.last_used_at,
            "client_disabled": r.disabled_at is not None,
        }
        for r in rows
    ]


async def list_all_grants(db: AsyncSession) -> list[dict]:
    """Every grant joined with its client + owner (for the admin view)."""
    from app.identity.person_models import IdentityPerson

    rows = (await db.execute(
        select(
            OAuthGrant.client_id, OAuthGrant.user_id, OAuthGrant.granted_at,
            OAuthGrant.last_used_at, OAuthGrant.revoked_at,
            OAuthGrant.revoked_by, OAuthClient.client_name,
            OAuthClient.disabled_at, OAuthClient.created_at,
            IdentityPerson.display_name, IdentityPerson.work_email,
        )
        .join(OAuthClient, OAuthClient.client_id == OAuthGrant.client_id)
        .join(IdentityPerson, IdentityPerson.id == OAuthGrant.user_id, isouter=True)
        .order_by(OAuthGrant.granted_at.desc())
    )).all()
    return [
        {
            "client_id": r.client_id,
            "client_name": r.client_name,
            "client_disabled": r.disabled_at is not None,
            "client_created_at": r.created_at,
            "user_id": str(r.user_id),
            "user_display_name": r.display_name,
            "user_email": r.work_email,
            "granted_at": r.granted_at,
            "last_used_at": r.last_used_at,
            "revoked_at": r.revoked_at,
            "revoked_by": r.revoked_by,
        }
        for r in rows
    ]


async def set_client_disabled(
    db: AsyncSession, *, client_id: str, disabled: bool,
) -> bool:
    """Admin kill-switch: flip a client's ``disabled_at``. Returns True if
    the client exists. A disabled client immediately fails authorize +
    every existing grant (middleware re-checks per request)."""
    row = (await db.execute(
        select(OAuthClient).where(OAuthClient.client_id == client_id)
    )).scalar_one_or_none()
    if row is None:
        return False
    row.disabled_at = datetime.now(UTC) if disabled else None
    await db.flush()
    return True


__all__ = [
    "AuthorizationCode",
    "OAuthError",
    "client_registration_payload",
    "consume_authorization_code",
    "consume_refresh_jti",
    "get_active_client",
    "is_grant_active",
    "issue_authorization_code",
    "list_all_grants",
    "list_user_grants",
    "record_grant",
    "register_client",
    "enforce_registration_quota",
    "prune_stale_ungranted_clients",
    "revoke_grant",
    "revoke_refresh_jti",
    "set_client_disabled",
    "store_refresh_jti",
    "touch_grant",
    "verify_pkce",
]