<?php
/**
 * JWT/JWK verification. Verify-only: a Company never signs anything in this
 * protocol, so there is no signing key to protect here.
 *
 * Built on the vendored firebase/php-jwt (ops/vendor-jwt.sh), narrowed to what
 * the spec and Guides allow:
 *   - algorithms ES256 and RS256 only; `none` and every HMAC alg refused
 *     before the library sees the token (Guides, "Signing and verifying JWTs");
 *   - keys come only from a verified source (the agent's jwks_uri, or a DPoP
 *     proof's own `jwk` header); `jku`, `x5u` and embedded keys are ignored;
 *   - a key set entry carrying private material is refused, not used;
 *   - RSA keys must be at least 2048 bits (the library enforces this too
 *     since 7.0.0; checked here as well so the refusal has our error code).
 *
 * @package AgentAdapt
 */

defined( 'ABSPATH' ) || exit;

use AgentAdapt\Vendor\Firebase\JWT\JWK;
use AgentAdapt\Vendor\Firebase\JWT\JWT;
use AgentAdapt\Vendor\Firebase\JWT\Key;

/**
 * JOSE helpers.
 */
final class AgentAdapt_Jose {

	public const ALGS = array( 'ES256', 'RS256' );

	/**
	 * Clock skew allowed on iat / nbf / exp (Guides suggest about 30 s).
	 */
	public const LEEWAY = 30;

	private const PRIVATE_MEMBERS = array( 'd', 'p', 'q', 'dp', 'dq', 'qi', 'oth', 'k' );

	/**
	 * Decode a JWT's header and payload WITHOUT verifying it. Only for reading
	 * routing hints (kid, typ, iss) before choosing a key. Never trust a value
	 * from here until verify() has returned the same token's payload.
	 *
	 * @param string $jwt Compact JWS.
	 * @return array{0: array<string,mixed>, 1: array<string,mixed>}|null
	 */
	public static function peek( string $jwt ): ?array {
		$parts = explode( '.', $jwt, 4 );
		if ( 3 !== count( $parts ) || strlen( $jwt ) > 16384 ) {
			return null;
		}
		try {
			$h = json_decode( JWT::urlsafeB64Decode( $parts[0] ), true, 8, JSON_THROW_ON_ERROR );
			$p = json_decode( JWT::urlsafeB64Decode( $parts[1] ), true, 8, JSON_THROW_ON_ERROR );
		} catch ( Throwable $e ) {
			return null;
		}
		return ( is_array( $h ) && is_array( $p ) ) ? array( $h, $p ) : null;
	}

	/**
	 * Build verification keys from a JWKS, keyed by kid ('' for a lone key
	 * without one). Unusable entries are skipped, never "fixed up".
	 *
	 * @param array<string,mixed> $jwks JWK Set.
	 * @return array<string, Key>
	 */
	public static function keys_from_jwks( array $jwks ): array {
		$out = array();
		foreach ( (array) ( $jwks['keys'] ?? array() ) as $jwk ) {
			if ( ! is_array( $jwk ) ) {
				continue;
			}
			if ( isset( $jwk['use'] ) && 'sig' !== $jwk['use'] ) {
				continue;
			}
			$key = self::key_from_jwk( $jwk );
			if ( $key instanceof Key ) {
				$out[ (string) ( $jwk['kid'] ?? '' ) ] = $key;
			}
		}
		return $out;
	}

	/**
	 * One public JWK → Key, or null when it is not an acceptable public
	 * signing key.
	 *
	 * @param array<string,mixed> $jwk JWK.
	 */
	public static function key_from_jwk( array $jwk ): ?Key {
		foreach ( self::PRIVATE_MEMBERS as $m ) {
			if ( array_key_exists( $m, $jwk ) ) {
				return null;
			}
		}
		$kty = (string) ( $jwk['kty'] ?? '' );
		$alg = (string) ( $jwk['alg'] ?? ( 'EC' === $kty ? 'ES256' : ( 'RSA' === $kty ? 'RS256' : '' ) ) );
		if ( ! in_array( $alg, self::ALGS, true ) ) {
			return null;
		}
		if ( ( 'ES256' === $alg && ( 'EC' !== $kty || 'P-256' !== ( $jwk['crv'] ?? '' ) ) ) || ( 'RS256' === $alg && 'RSA' !== $kty ) ) {
			return null;
		}
		try {
			$key = JWK::parseKey( $jwk, $alg );
		} catch ( Throwable $e ) {
			return null;
		}
		if ( ! $key instanceof Key ) {
			return null;
		}
		if ( 'RS256' === $alg ) {
			$details = openssl_pkey_get_details( openssl_pkey_get_public( $key->getKeyMaterial() ) );
			if ( ! is_array( $details ) || (int) ( $details['bits'] ?? 0 ) < 2048 ) {
				return null;
			}
		}
		return $key;
	}

	/**
	 * Verify a compact JWS against a set of keys and return its claims.
	 *
	 * The key is chosen by the header's `kid`; a header without one is
	 * accepted only when the set holds exactly one key. The header `alg` must
	 * be ES256/RS256 AND match the chosen key's algorithm (the library checks
	 * the latter; this checks the former first so HMAC never reaches it).
	 *
	 * @param string             $jwt  Compact JWS.
	 * @param array<string, Key> $keys kid → Key.
	 * @param string|null        $typ  Required `typ` header, case-insensitive, or null.
	 * @return array{header: array<string,mixed>, claims: array<string,mixed>}|WP_Error
	 */
	public static function verify( string $jwt, array $keys, ?string $typ = null ) {
		$peek = self::peek( $jwt );
		if ( null === $peek ) {
			return new WP_Error( 'malformed', 'not a compact JWS' );
		}
		[ $header ] = $peek;
		if ( ! in_array( $header['alg'] ?? '', self::ALGS, true ) ) {
			return new WP_Error( 'alg', 'algorithm not allowed' );
		}
		if ( null !== $typ && 0 !== strcasecmp( (string) ( $header['typ'] ?? '' ), $typ ) ) {
			return new WP_Error( 'typ', 'wrong typ' );
		}
		if ( ! $keys ) {
			return new WP_Error( 'no_key', 'no usable verification key' );
		}
		$kid = isset( $header['kid'] ) ? (string) $header['kid'] : null;
		if ( null === $kid ) {
			if ( 1 !== count( $keys ) ) {
				return new WP_Error( 'kid', 'kid required when the key set has several keys' );
			}
			$key = reset( $keys );
		} elseif ( isset( $keys[ $kid ] ) ) {
			$key = $keys[ $kid ];
		} else {
			return new WP_Error( 'unknown_kid', 'unknown kid' );
		}

		$prev         = JWT::$leeway;
		JWT::$leeway  = self::LEEWAY;
		try {
			$claims = JWT::decode( $jwt, $key );
		} catch ( Throwable $e ) {
			return new WP_Error( 'signature', 'verification failed' );
		} finally {
			JWT::$leeway = $prev;
		}
		return array(
			'header' => $header,
			'claims' => json_decode( (string) wp_json_encode( $claims ), true ),
		);
	}

	/**
	 * RFC 7638 thumbprint of a public JWK (EC or RSA): the DPoP `jkt`.
	 *
	 * @param array<string,mixed> $jwk Public JWK.
	 */
	public static function thumbprint( array $jwk ): string {
		$kty = (string) ( $jwk['kty'] ?? '' );
		if ( 'EC' === $kty ) {
			$members = array( 'crv' => $jwk['crv'] ?? '', 'kty' => 'EC', 'x' => $jwk['x'] ?? '', 'y' => $jwk['y'] ?? '' );
		} elseif ( 'RSA' === $kty ) {
			$members = array( 'e' => $jwk['e'] ?? '', 'kty' => 'RSA', 'n' => $jwk['n'] ?? '' );
		} else {
			return '';
		}
		return JWT::urlsafeB64Encode( hash( 'sha256', (string) wp_json_encode( $members, JSON_UNESCAPED_SLASHES ), true ) );
	}

	/**
	 * base64url(SHA-256(token)): the DPoP `ath` for an access token.
	 *
	 * @param string $token Access token.
	 */
	public static function ath( string $token ): string {
		return JWT::urlsafeB64Encode( hash( 'sha256', $token, true ) );
	}
}