<?php
/**
 * The token endpoint (spec 4.2; 4.4–4.8 arrive with sign-in in Phase 2).
 *
 * Grants:
 *   - JWT bearer (spec 4.2): start a signed-out Session, or renew one;
 *   - authorization_code (Direct Sign-In, spec 4.5) and device_code (Device
 *     Sign-In, spec 4.6): sign the given Session in and issue an Account Token;
 *   - refresh_token (the Account Token, spec 4.8): a signed-in Session Token
 *     for an existing or a new Session, optionally with fewer scopes.
 * A request carries the client's private_key_jwt assertion AND a DPoP proof,
 * and the issued Session Token is bound to the proof's key. The one exception
 * is spec 4.3's Bearer token for an MCP server: requested with `resource` set
 * to that server, and good there alone. The server doesn't support DPoP, so the
 * token is Bearer even if the request carries a proof (Version 0.2). Free core serves
 * no MCP server, so it never issues one. A public client (MCP's own sign-in,
 * where a stock MCP client has no keys) may only redeem a PKCE-protected
 * code or its refresh token, and only for such a resource.
 *
 * Check order matters for what an outsider can make us do: the request shape
 * and per-address rate limit come first (cheap), then client authentication
 * (may fetch the agent's metadata, itself rate-limited for unknown agents),
 * then DPoP, then the grant.
 *
 * @package AgentAdapt
 */

defined( 'ABSPATH' ) || exit;

/**
 * Token endpoint.
 */
final class AgentAdapt_Token_Endpoint {

	public const GRANT_JWT_BEARER = 'urn:ietf:params:oauth:grant-type:jwt-bearer';
	public const GRANT_DEVICE     = 'urn:ietf:params:oauth:grant-type:device_code';

	/**
	 * Requests per minute per remote address, and Session starts per hour per
	 * agent and per subject. Filterable via agentadapt_rate_limit.
	 */
	private const PER_IP_PER_MINUTE        = 120;
	private const SESSIONS_PER_AGENT_HOUR  = 3000;
	private const SESSIONS_PER_SUBJECT_HOUR = 120;

	/**
	 * POST /agent-adapt/oauth/token
	 */
	public static function handle(): void {
		if ( 'POST' !== AgentAdapt_Http::method() ) {
			AgentAdapt_Http::json( 405, array( 'error' => 'invalid_request', 'error_description' => 'POST only' ), array( 'Allow' => 'POST', 'Cache-Control' => 'no-store' ) );
		}
		if ( AgentAdapt_Http::token_in_query() ) {
			AgentAdapt_Http::oauth_error( 400, 'invalid_request', 'credentials are never accepted in the URL' );
		}
		if ( ! str_starts_with( strtolower( (string) AgentAdapt_Http::header( 'Content-Type' ) ), 'application/x-www-form-urlencoded' ) ) {
			AgentAdapt_Http::oauth_error( 400, 'invalid_request', 'the body must be application/x-www-form-urlencoded' );
		}
		self::limit( 'token-ip:' . AgentAdapt_Http::client_ip(), self::PER_IP_PER_MINUTE, MINUTE_IN_SECONDS );

		$form     = AgentAdapt_Http::form();
		$endpoint = AgentAdapt_Urls::token_endpoint();
		$grant    = (string) ( $form['grant_type'] ?? '' );

		$agent = AgentAdapt_Client_Auth::authenticate( $form );
		if ( is_wp_error( $agent ) ) {
			if ( 'rate_limited' === $agent->get_error_code() ) {
				AgentAdapt_Http::rate_limited( 3600 );
			}
			AgentAdapt_Http::oauth_error( 401, 'invalid_client', $agent->get_error_message() );
		}

		$resource = trim( (string) ( $form['resource'] ?? '' ) );
		if ( '' !== $resource && ! AgentAdapt_Urls::known_resource( $resource ) ) {
			AgentAdapt_Http::oauth_error( 400, 'invalid_target', 'unknown resource' );
		}
		$public = ! empty( $agent['public'] );
		if ( $public && ! in_array( $grant, array( 'authorization_code', 'refresh_token' ), true ) ) {
			AgentAdapt_Http::oauth_error( 400, 'unauthorized_client', 'a public client can use only authorization_code and refresh_token' );
		}
		if ( '' !== $resource && AgentAdapt_Urls::is_bearer_resource( $resource ) ) {
			// Spec 4.3 (Version 0.2): this store's MCP server doesn't support DPoP, so a token
			// for it is a Bearer token even when the request carries a proof, never a DPoP one.
			$bound = array( 'jkt' => '', 'resource' => $resource );
		} else {
			if ( $public ) {
				AgentAdapt_Http::oauth_error( 400, 'invalid_target', 'a public client can get tokens only for this store\'s MCP server: send it as resource' );
			}
			$dpop = AgentAdapt_Dpop::verify_request( 'POST', $endpoint );
			if ( is_wp_error( $dpop ) ) {
				AgentAdapt_Http::oauth_error( 400, 'invalid_dpop_proof', $dpop->get_error_message() );
			}
			$bound = array( 'jkt' => (string) $dpop['jkt'], 'resource' => '' );
		}

		switch ( $grant ) {
			case self::GRANT_JWT_BEARER:
				self::jwt_bearer( $agent, $form, $bound );
				break;
			case 'authorization_code':
				if ( AgentAdapt_Signin::offers( 'direct' ) ) {
					self::authorization_code( $agent, $form, $bound );
				}
				break;
			case self::GRANT_DEVICE:
				if ( AgentAdapt_Signin::offers( 'device' ) ) {
					self::device_code( $agent, $form, $bound );
				}
				break;
			case 'refresh_token':
				self::refresh_token( $agent, $form, $bound );
				break;
		}
		AgentAdapt_Http::oauth_error( 400, 'unsupported_grant_type', 'this store does not offer that grant' );
	}

	/**
	 * authorization_code: Direct Sign-In's code exchange (spec 4.5).
	 *
	 * @param array<string,mixed>  $agent Resolved agent.
	 * @param array<string,string> $form  Form parameters.
	 * @param array<string,string>  $bound What the token binds to: jkt, or resource for Bearer.
	 */
	private static function authorization_code( array $agent, array $form, array $bound ): void {
		$sid = (string) ( $form['session_id'] ?? '' );
		// A public client is a stock MCP client that knows nothing of Sessions:
		// its sign-in starts a new one (spec §6). Everyone else names theirs.
		$session = ( '' === $sid && ! empty( $agent['public'] ) ) ? null : self::agent_session( $agent, $sid );
		$code    = AgentAdapt_Codes::consume_auth( (string) ( $form['code'] ?? '' ), $agent['client_hash'] );
		if ( is_wp_error( $code ) ) {
			AgentAdapt_Http::oauth_error( 400, 'invalid_grant', $code->get_error_message() );
		}
		if ( ( $form['redirect_uri'] ?? null ) !== $code['redirect_uri'] ) {
			AgentAdapt_Http::oauth_error( 400, 'invalid_grant', 'redirect_uri does not match the authorization request' );
		}
		$verifier = (string) ( $form['code_verifier'] ?? '' );
		$expected = rtrim( strtr( base64_encode( hash( 'sha256', $verifier, true ) ), '+/', '-_' ), '=' ); // phpcs:ignore WordPress.PHP.DiscouragedPHPFunctions.obfuscation_base64_encode -- PKCE S256 (RFC 7636 §4.2).
		if ( ! preg_match( '/^[A-Za-z0-9._~-]{43,128}$/', $verifier ) || ! hash_equals( (string) $code['code_challenge'], $expected ) ) {
			AgentAdapt_Http::oauth_error( 400, 'invalid_grant', 'code_verifier does not match' );
		}
		// RFC 8707 §2.2: a code issued for a resource is redeemed for that resource.
		if ( '' !== (string) ( $code['resource'] ?? '' ) && ! AgentAdapt_Urls::same_resource( (string) $code['resource'], (string) ( $form['resource'] ?? '' ) ) ) {
			AgentAdapt_Http::oauth_error( 400, 'invalid_target', 'resource does not match the authorization request' );
		}
		if ( null === $session ) {
			$subject = AgentAdapt_Sessions::subject_id( $agent['client_id'], AgentAdapt_Sessions::account_sub( (int) $code['user_id'] ) );
			$session = AgentAdapt_Sessions::start( $agent['client_hash'], $subject );
			AgentAdapt_Events::record( 'session.started', null, array( 'client_hash' => $agent['client_hash'], 'subject_id' => $subject, 'session_id' => $session['session_id'], 'channel' => 'protocol' ) );
		}
		self::sign_in( $agent, $session, (int) $code['user_id'], (string) $code['granted_scope'], $bound, 'direct', (string) $code['code_hash'] );
	}

	/**
	 * device_code: Device Sign-In's poll (spec 4.6).
	 *
	 * @param array<string,mixed>  $agent Resolved agent.
	 * @param array<string,string> $form  Form parameters.
	 * @param array<string,string>  $bound What the token binds to.
	 */
	private static function device_code( array $agent, array $form, array $bound ): void {
		$session = self::agent_session( $agent, (string) ( $form['session_id'] ?? '' ) );
		$row     = AgentAdapt_Codes::poll_device( (string) ( $form['device_code'] ?? '' ), $agent['client_hash'] );
		if ( is_wp_error( $row ) ) {
			AgentAdapt_Http::oauth_error( 400, $row->get_error_code(), $row->get_error_message() );
		}
		self::sign_in( $agent, $session, (int) $row['user_id'], (string) $row['granted_scope'], $bound, 'device' );
	}

	/**
	 * refresh_token: use the Account Token (spec 4.8).
	 *
	 * @param array<string,mixed>  $agent Resolved agent.
	 * @param array<string,string> $form  Form parameters.
	 * @param array<string,string>  $bound What the token binds to.
	 */
	private static function refresh_token( array $agent, array $form, array $bound ): void {
		$grant = AgentAdapt_Grants::find( (string) ( $form['refresh_token'] ?? '' ), $agent['client_hash'] );
		if ( null === $grant ) {
			AgentAdapt_Http::oauth_error( 400, 'invalid_grant', 'the Account Token is unknown, expired or revoked' );
		}
		$user_id = (int) $grant['user_id'];
		if ( ! AgentAdapt_Eligibility::eligible( $user_id ) ) {
			AgentAdapt_Grants::revoke( (string) $grant['grant_id'], 'no_longer_eligible' );
			AgentAdapt_Http::oauth_error( 400, 'invalid_grant', 'the Account Token is unknown, expired or revoked' );
		}
		$have  = array_values( array_filter( explode( ' ', (string) $grant['scope'] ) ) );
		$asked = isset( $form['scope'] ) && '' !== trim( (string) $form['scope'] ) ? array_values( array_filter( explode( ' ', trim( (string) $form['scope'] ) ) ) ) : $have;
		if ( array_diff( $asked, $have ) ) {
			// May narrow, never widen (spec 4.8).
			AgentAdapt_Http::oauth_error( 400, 'invalid_scope', 'an Account Token cannot grant more than the customer approved' );
		}
		if ( isset( $form['session_id'] ) && '' !== $form['session_id'] ) {
			$session = AgentAdapt_Sessions::get( (string) $form['session_id'] );
			if ( ! $session || $session['client_hash'] !== $agent['client_hash'] || (int) $session['subject_id'] !== (int) $grant['subject_id'] || ( ! empty( $agent['public'] ) && AgentAdapt_Sessions::end_if_signed_out( $session ) ) ) {
				AgentAdapt_Http::oauth_error( 400, 'invalid_session', 'start a new Session' );
			}
		} else {
			$session = AgentAdapt_Sessions::start( $agent['client_hash'], (int) $grant['subject_id'] );
			AgentAdapt_Events::record( 'session.started', null, array( 'client_hash' => $agent['client_hash'], 'subject_id' => (int) $grant['subject_id'], 'session_id' => $session['session_id'], 'channel' => 'protocol' ) );
		}
		$ok = AgentAdapt_Sessions::sign_in( $session, $user_id, (string) $grant['grant_id'] );
		if ( is_wp_error( $ok ) ) {
			AgentAdapt_Http::oauth_error( 400, 'account_mismatch', $ok->get_error_message() );
		}
		AgentAdapt_Grants::touch( (string) $grant['grant_id'] );
		$issued = AgentAdapt_Tokens::issue_session_token( $session['session_id'], implode( ' ', $asked ), $bound['jkt'], $user_id, $bound['resource'] );
		self::respond( $issued, implode( ' ', $asked ), $session['session_id'], true, '' === $bound['jkt'] );
	}

	/**
	 * Finish Direct or Device Sign-In and answer with the token response.
	 *
	 * @param array<string,mixed>  $agent     Agent.
	 * @param array<string,mixed>  $session   The Session being signed in.
	 * @param int                  $user_id   Account that approved.
	 * @param string               $scope     Scopes the customer approved.
	 * @param array<string,string> $bound     What the token binds to.
	 * @param string               $type      Sign-in type, for the event trail.
	 * @param string               $code_hash Authorization code (for reuse revocation), or ''.
	 */
	private static function sign_in( array $agent, array $session, int $user_id, string $scope, array $bound, string $type, string $code_hash = '' ): void {
		$body = self::complete_sign_in( $agent, $session, $user_id, $scope, $bound, $type, $code_hash );
		if ( is_wp_error( $body ) ) {
			AgentAdapt_Http::oauth_error( 400, $body->get_error_code(), $body->get_error_message() );
		}
		AgentAdapt_Http::json( 200, $body, array( 'Cache-Control' => 'no-store', 'Pragma' => 'no-cache' ) );
	}

	/**
	 * The one way a sign-in ends, whatever its type (spec 4.4, "All three
	 * types end the same way"): the account is checked again, the Session's
	 * account binding is honoured, then the Account Token is issued, the
	 * Session is signed in, a signed-in Session Token is issued, and the
	 * sign-in is recorded. Premium sign-in types finish through here too, so
	 * Connected agents, revocation and Account Token use treat every sign-in
	 * alike.
	 *
	 * @param array<string,mixed>  $agent     Agent: at least client_hash.
	 * @param array<string,mixed>  $session   The Session being signed in (AgentAdapt_Sessions::get()).
	 * @param int                  $user_id   The account.
	 * @param string               $scope     Scopes granted, space-separated.
	 * @param array<string,string> $bound     What the token binds to: jkt, or resource for Bearer.
	 * @param string               $type      Sign-in type (`direct`, `device`, or one a premium module adds), recorded with the event.
	 * @param string               $code_hash Authorization code (for reuse revocation), or ''.
	 * @return array<string,mixed>|WP_Error The token response body (spec 4.4), or invalid_grant (the account can't be connected) or account_mismatch.
	 */
	public static function complete_sign_in( array $agent, array $session, int $user_id, string $scope, array $bound, string $type, string $code_hash = '' ) {
		if ( ! AgentAdapt_Eligibility::eligible( $user_id ) ) {
			return new WP_Error( 'invalid_grant', 'this account cannot be connected' );
		}
		$bound_user = (int) $session['bound_user_id'];
		if ( 0 !== $bound_user && $bound_user !== $user_id ) {
			return new WP_Error( 'account_mismatch', 'this Session belongs to a different account; start a new Session' );
		}
		$grant = AgentAdapt_Grants::issue( (string) $agent['client_hash'], (int) $session['subject_id'], $user_id, $scope );
		if ( '' !== $code_hash ) {
			AgentAdapt_Codes::link_grant( $code_hash, $grant['grant_id'] );
		}
		AgentAdapt_Sessions::sign_in( $session, $user_id, $grant['grant_id'] );
		$issued = AgentAdapt_Tokens::issue_session_token( (string) $session['session_id'], $scope, (string) $bound['jkt'], $user_id, (string) ( $bound['resource'] ?? '' ) );
		AgentAdapt_Events::record( 'session.signed_in', null, array( 'client_hash' => (string) $agent['client_hash'], 'subject_id' => (int) $session['subject_id'], 'session_id' => (string) $session['session_id'], 'channel' => 'protocol', 'meta' => array( 'scope' => $scope, 'type' => $type ) ) );
		return self::body( $issued, $scope, (string) $session['session_id'], true, $grant, '' === (string) $bound['jkt'] );
	}

	/**
	 * The Session a sign-in completes: it must be live and this agent's.
	 *
	 * @param array<string,mixed> $agent      Agent.
	 * @param string              $session_id Session ID from the request.
	 * @return array<string,mixed>
	 */
	private static function agent_session( array $agent, string $session_id ): array {
		$session = AgentAdapt_Sessions::get( $session_id );
		// A public client's Session that lost its sign-in has ended (spec §6).
		if ( ! $session || $session['client_hash'] !== $agent['client_hash'] || ( ! empty( $agent['public'] ) && AgentAdapt_Sessions::end_if_signed_out( $session ) ) ) {
			AgentAdapt_Http::oauth_error( 400, 'invalid_session', 'session_id must name a live Session of this agent' );
		}
		return $session;
	}

	/**
	 * The token response body (spec 4.2 / 4.4).
	 *
	 * @param array{token: string, expires_in: int}                        $issued     Session Token.
	 * @param string                                                       $scope      Granted scopes.
	 * @param string                                                       $session_id Session.
	 * @param bool                                                         $signed_in  Signed in.
	 * @param array{token: string, grant_id: string, expires_in: int}|null $grant      New Account Token, if any.
	 * @param bool                                                         $bearer     A Bearer token (spec 4.3).
	 * @return array<string,mixed>
	 */
	private static function body( array $issued, string $scope, string $session_id, bool $signed_in, ?array $grant = null, bool $bearer = false ): array {
		$body = array(
			'access_token' => $issued['token'],
			'token_type'   => $bearer ? 'Bearer' : 'DPoP',
			'expires_in'   => $issued['expires_in'],
			'scope'        => $scope,
			'session_id'   => $session_id,
			'signed_in'    => $signed_in,
		);
		if ( $grant ) {
			$body['refresh_token']            = $grant['token'];
			$body['refresh_token_expires_in'] = $grant['expires_in'];
		}
		return $body;
	}

	/**
	 * Answer with the token response (spec 4.2 / 4.4).
	 *
	 * @param array{token: string, expires_in: int} $issued     Session Token.
	 * @param string                                $scope      Granted scopes.
	 * @param string                                $session_id Session.
	 * @param bool                                  $signed_in  Signed in.
	 * @param bool                                  $bearer     A Bearer token (spec 4.3).
	 */
	private static function respond( array $issued, string $scope, string $session_id, bool $signed_in, bool $bearer = false ): void {
		AgentAdapt_Http::json( 200, self::body( $issued, $scope, $session_id, $signed_in, null, $bearer ), array( 'Cache-Control' => 'no-store', 'Pragma' => 'no-cache' ) );
	}

	/**
	 * JWT bearer grant: start or renew a signed-out Session.
	 *
	 * @param array<string,mixed>  $agent Resolved agent.
	 * @param array<string,string> $form  Form parameters.
	 * @param array<string,string>  $bound What the token binds to.
	 */
	private static function jwt_bearer( array $agent, array $form, array $bound ): void {
		$claims = AgentAdapt_Client_Auth::verify_agent_jwt( $agent, (string) ( $form['assertion'] ?? '' ), 'assertion' );
		if ( is_wp_error( $claims ) ) {
			AgentAdapt_Http::oauth_error( 400, 'invalid_grant', 'assertion: ' . $claims->get_error_message() );
		}
		// Spec 4.2: aud is this store's issuer as a single string, like the
		// client assertion's (4.1), so an assertion minted for one Company is
		// useless at another, including one that names our token endpoint.
		if ( ! AgentAdapt_Client_Auth::aud_is_issuer( $claims ) ) {
			AgentAdapt_Http::oauth_error( 400, 'invalid_grant', 'assertion aud must be this store\'s issuer (auth.issuer), as a single string' );
		}
		$sub = $claims['sub'] ?? null;
		if ( ! is_string( $sub ) || '' === $sub || strlen( $sub ) > 255 ) {
			AgentAdapt_Http::oauth_error( 400, 'invalid_grant', 'assertion sub must be the User ID' );
		}
		if ( '' !== trim( (string) ( $form['scope'] ?? '' ) ) ) {
			// A signed-out Session has no account scopes to narrow (spec 4.4).
			AgentAdapt_Http::oauth_error( 400, 'invalid_scope', 'a signed-out Session has no scopes' );
		}

		$subject_id = AgentAdapt_Sessions::subject_id( $agent['client_id'], $sub );

		if ( isset( $form['session_id'] ) ) {
			$session = AgentAdapt_Sessions::get( (string) $form['session_id'] );
			if ( ! $session || $session['client_hash'] !== $agent['client_hash'] || (int) $session['subject_id'] !== $subject_id ) {
				// Unknown, ended, or another agent's / User's: one answer for all.
				AgentAdapt_Http::oauth_error( 400, 'invalid_session', 'start a new Session' );
			}
			AgentAdapt_Sessions::touch( $session );
			$event = 'session.renewed';
		} else {
			self::limit( 'sessions-agent:' . $agent['client_hash'], self::SESSIONS_PER_AGENT_HOUR, HOUR_IN_SECONDS );
			self::limit( 'sessions-subject:' . $subject_id, self::SESSIONS_PER_SUBJECT_HOUR, HOUR_IN_SECONDS );
			$session = AgentAdapt_Sessions::start( $agent['client_hash'], $subject_id );
			$event   = 'session.started';
		}

		$issued = AgentAdapt_Tokens::issue_session_token( $session['session_id'], '', $bound['jkt'], 0, $bound['resource'] );
		AgentAdapt_Events::record(
			$event,
			null,
			array( 'client_hash' => $agent['client_hash'], 'subject_id' => $subject_id, 'session_id' => $session['session_id'], 'channel' => 'protocol' )
		);
		AgentAdapt_Http::json(
			200,
			array(
				'access_token' => $issued['token'],
				'token_type'   => '' === $bound['jkt'] ? 'Bearer' : 'DPoP',
				'expires_in'   => $issued['expires_in'],
				'scope'        => '',
				'session_id'   => $session['session_id'],
				'signed_in'    => false,
			),
			array( 'Cache-Control' => 'no-store', 'Pragma' => 'no-cache' )
		);
	}

	/**
	 * Count a hit; answer 429 with Retry-After when over.
	 *
	 * @param string $bucket Bucket.
	 * @param int    $limit  Limit.
	 * @param int    $window Window seconds.
	 */
	private static function limit( string $bucket, int $limit, int $window ): void {
		$r = AgentAdapt_Rate_Limit::hit( $bucket, $limit, $window );
		if ( ! $r['allowed'] ) {
			AgentAdapt_Http::rate_limited( (int) $r['retry_after'] );
		}
	}
}