<?php
/**
 * OAuth 2.1 server orchestrator for the MCP endpoint.
 *
 * Part of AI-9. Phase 1 implemented here: discovery metadata (RFC 8414 + RFC 9728),
 * `.well-known` rewrite endpoints, and the 401 `WWW-Authenticate` challenge that lets
 * MCP clients (Claude.ai, Claude Desktop, Cursor) auto-discover the auth server.
 *
 * Authorization / token / registration endpoints (RFC 7591/7636) are added in later
 * phases. The whole subsystem is gated behind `awfah_oauth_settings['enabled']`
 * (default false) until the full flow ships and passes a security review.
 *
 * @package AtlasAiDev\Awfah
 * @subpackage OAuth
 */

declare(strict_types=1);

namespace AtlasAiDev\Awfah\OAuth;

defined( 'ABSPATH' ) || exit;

/**
 * Class OAuth_Server
 */
class OAuth_Server {

	const SETTINGS_OPTION  = 'awfah_oauth_settings';
	const REWRITE_OPTION   = 'awfah_oauth_rewrite_version';
	const REWRITE_VERSION  = '4';
	const REST_NAMESPACE   = 'awfah_oauth/v1';
	const QUERY_VAR        = 'awfah_well_known';

	/**
	 * Wire hooks.
	 */
	public function __construct() {
		add_action( 'init', array( $this, 'on_init' ) );
		add_filter( 'query_vars', array( $this, 'register_query_var' ) );
		add_action( 'template_redirect', array( $this, 'maybe_serve_well_known' ) );
		add_action( 'rest_api_init', array( $this, 'register_metadata_routes' ) );
		add_filter( 'rest_post_dispatch', array( $this, 'add_challenge_header' ), 10, 3 );
		// Don't let WP append a canonical trailing slash to the .well-known URLs.
		add_filter( 'redirect_canonical', array( $this, 'prevent_well_known_redirect' ) );

		// Sub-controllers (each self-hooks).
		new Registration_Controller();
		new Authorize_Controller();
		new Token_Controller();
		new Token_Validator();
	}

	/**
	 * Generate a random opaque token with a type prefix.
	 *
	 * @param string $prefix Prefix (e.g. 'access_', 'refresh_', 'code_', 'client_').
	 * @return string
	 */
	public static function random_token( string $prefix ): string {
		return $prefix . bin2hex( random_bytes( 32 ) );
	}

	/**
	 * Best-effort client IP for rate limiting / audit.
	 *
	 * @return string
	 */
	public static function client_ip(): string {
		return isset( $_SERVER['REMOTE_ADDR'] ) ? sanitize_text_field( wp_unslash( $_SERVER['REMOTE_ADDR'] ) ) : '';
	}

	/**
	 * Read the Authorization header across server configurations.
	 *
	 * @return string
	 */
	public static function get_auth_header(): string {
		$auth = '';
		if ( isset( $_SERVER['HTTP_AUTHORIZATION'] ) ) {
			$auth = wp_unslash( $_SERVER['HTTP_AUTHORIZATION'] ); // phpcs:ignore WordPress.Security.ValidatedSanitizedInput
		} elseif ( isset( $_SERVER['REDIRECT_HTTP_AUTHORIZATION'] ) ) {
			$auth = wp_unslash( $_SERVER['REDIRECT_HTTP_AUTHORIZATION'] ); // phpcs:ignore WordPress.Security.ValidatedSanitizedInput
		} elseif ( function_exists( 'apache_request_headers' ) ) {
			foreach ( (array) apache_request_headers() as $k => $v ) {
				if ( 0 === strcasecmp( (string) $k, 'authorization' ) ) {
					$auth = $v;
					break;
				}
			}
		}
		return is_string( $auth ) ? trim( $auth ) : '';
	}

	/* --------------------------------------------------------------------- *
	 * Settings / helpers
	 * --------------------------------------------------------------------- */

	/**
	 * Whether the OAuth subsystem is enabled.
	 *
	 * @return bool
	 */
	public static function is_enabled(): bool {
		$settings = (array) get_option( self::SETTINGS_OPTION, array() );
		$enabled  = ! empty( $settings['enabled'] );

		/**
		 * Filter whether the OAuth 2.1 server is enabled.
		 *
		 * @param bool $enabled
		 */
		return (bool) apply_filters( 'awfah_oauth_enabled', $enabled );
	}

	/**
	 * OAuth scopes this server understands.
	 *
	 * @return string[]
	 */
	public static function scopes(): array {
		return (array) apply_filters( 'awfah_oauth_scopes', array( 'mcp:read', 'mcp:write', 'mcp:woocommerce' ) );
	}

	/**
	 * The OAuth issuer (the WordPress site).
	 *
	 * @return string
	 */
	public static function issuer(): string {
		return untrailingslashit( home_url() );
	}

	/**
	 * The protected resource URL (the MCP endpoint).
	 *
	 * @return string
	 */
	public static function resource_url(): string {
		$ns    = 'awfah_mcp';
		$route = 'mcp';
		if ( class_exists( '\AtlasAiDev\Awfah\Plugin' ) && method_exists( '\AtlasAiDev\Awfah\Plugin', 'get_mcp_settings' ) ) {
			$settings = \AtlasAiDev\Awfah\Plugin::get_mcp_settings();
			$ns       = $settings['namespace'] ?? $ns;
			$route    = $settings['route'] ?? $route;
		}
		return rest_url( $ns . '/' . $route );
	}

	/**
	 * URL of the protected-resource metadata document.
	 *
	 * @return string
	 */
	public static function protected_resource_metadata_url(): string {
		return home_url( '/.well-known/oauth-protected-resource' );
	}

	/* --------------------------------------------------------------------- *
	 * Rewrites + .well-known
	 * --------------------------------------------------------------------- */

	/**
	 * Register rewrite rules, ensure tables, flush once when needed.
	 *
	 * @return void
	 */
	public function on_init(): void {
		add_rewrite_rule( '^\.well-known/oauth-authorization-server/?$', 'index.php?' . self::QUERY_VAR . '=as', 'top' );
		// Match the plain PRM path AND the RFC 9728 path-inserted form
		// (.well-known/oauth-protected-resource/<resource-path>), which MCP clients
		// request during the token-exchange phase when they have no WWW-Authenticate hint.
		add_rewrite_rule( '^\.well-known/oauth-protected-resource(/.*)?$', 'index.php?' . self::QUERY_VAR . '=prm', 'top' );
		// OIDC-style discovery alias. Some MCP clients (e.g. mcp-remote) only probe
		// {path}/.well-known/openid-configuration for subdirectory/pathful issuers,
		// never {path}/.well-known/oauth-authorization-server — so serve it here too.
		add_rewrite_rule( '^\.well-known/openid-configuration/?$', 'index.php?' . self::QUERY_VAR . '=oidc', 'top' );
		add_rewrite_rule( '^awfah-oauth/authorize/?$', 'index.php?' . Authorize_Controller::QUERY_VAR . '=1', 'top' );

		if ( (string) get_option( self::REWRITE_OPTION ) !== self::REWRITE_VERSION ) {
			OAuth_Repository::maybe_create_tables();
			flush_rewrite_rules( false );
			update_option( self::REWRITE_OPTION, self::REWRITE_VERSION, false );
		}
	}

	/**
	 * Prevent the canonical trailing-slash redirect on the .well-known endpoints.
	 *
	 * @param string|false $redirect_url The canonical URL WP wants to redirect to.
	 * @return string|false
	 */
	public function prevent_well_known_redirect( $redirect_url ) {
		if ( '' !== (string) get_query_var( self::QUERY_VAR ) ) {
			return false;
		}
		return $redirect_url;
	}

	/**
	 * Register the well-known query var.
	 *
	 * @param array $vars Query vars.
	 * @return array
	 */
	public function register_query_var( $vars ) {
		$vars[] = self::QUERY_VAR;
		$vars[] = Authorize_Controller::QUERY_VAR;
		return $vars;
	}

	/**
	 * Serve the discovery metadata documents at the .well-known paths.
	 *
	 * @return void
	 */
	public function maybe_serve_well_known(): void {
		$which = get_query_var( self::QUERY_VAR );
		if ( '' === $which || null === $which ) {
			return;
		}

		if ( ! self::is_enabled() ) {
			status_header( 404 );
			exit;
		}

		// 'prm' => protected-resource metadata; 'as' and 'oidc' => authorization-server metadata.
		$data = 'prm' === $which ? self::protected_resource_metadata() : self::authorization_server_metadata();

		header( 'Content-Type: application/json; charset=utf-8' );
		header( 'Access-Control-Allow-Origin: *' );
		header( 'Cache-Control: public, max-age=3600' );
		echo wp_json_encode( $data );
		exit;
	}

	/* --------------------------------------------------------------------- *
	 * REST metadata routes (mirror of .well-known, for clients that probe REST)
	 * --------------------------------------------------------------------- */

	/**
	 * Register REST metadata routes.
	 *
	 * @return void
	 */
	public function register_metadata_routes(): void {
		register_rest_route(
			self::REST_NAMESPACE,
			'/protected-resource-metadata',
			array(
				'methods'             => 'GET',
				'permission_callback' => '__return_true',
				'callback'            => static function () {
					if ( ! self::is_enabled() ) {
						return new \WP_REST_Response( array( 'error' => 'oauth_disabled' ), 404 );
					}
					return new \WP_REST_Response( self::protected_resource_metadata(), 200 );
				},
			)
		);
		register_rest_route(
			self::REST_NAMESPACE,
			'/authorization-server-metadata',
			array(
				'methods'             => 'GET',
				'permission_callback' => '__return_true',
				'callback'            => static function () {
					if ( ! self::is_enabled() ) {
						return new \WP_REST_Response( array( 'error' => 'oauth_disabled' ), 404 );
					}
					return new \WP_REST_Response( self::authorization_server_metadata(), 200 );
				},
			)
		);
		// Minimal JWKS endpoint. We issue opaque tokens (no JWT/id_token), so the
		// key set is empty — present only because the OIDC discovery doc references it.
		register_rest_route(
			self::REST_NAMESPACE,
			'/jwks',
			array(
				'methods'             => 'GET',
				'permission_callback' => '__return_true',
				'callback'            => static function () {
					return new \WP_REST_Response( array( 'keys' => array() ), 200 );
				},
			)
		);
	}

	/* --------------------------------------------------------------------- *
	 * Metadata documents
	 * --------------------------------------------------------------------- */

	/**
	 * RFC 8414 Authorization Server Metadata.
	 *
	 * @return array
	 */
	public static function authorization_server_metadata(): array {
		$base = rest_url( self::REST_NAMESPACE );
		return array(
			'issuer'                                => self::issuer(),
			'authorization_endpoint'                => home_url( '/awfah-oauth/authorize' ),
			'token_endpoint'                        => $base . '/token',
			'registration_endpoint'                 => $base . '/register',
			'revocation_endpoint'                   => $base . '/revoke',
			'scopes_supported'                      => self::scopes(),
			'response_types_supported'              => array( 'code' ),
			'grant_types_supported'                 => array( 'authorization_code', 'refresh_token' ),
			'code_challenge_methods_supported'      => array( 'S256' ),
			'token_endpoint_auth_methods_supported' => array( 'none', 'client_secret_post', 'client_secret_basic' ),
			// OIDC-discovery-required fields. We are an OAuth 2.1 server (opaque
			// tokens, no id_tokens), but clients that fetch this doc via
			// /.well-known/openid-configuration validate it against the OIDC schema
			// and require these keys. They are inert for the OAuth code flow.
			'jwks_uri'                              => $base . '/jwks',
			'subject_types_supported'               => array( 'public' ),
			'id_token_signing_alg_values_supported' => array( 'RS256' ),
		);
	}

	/**
	 * RFC 9728 Protected Resource Metadata.
	 *
	 * @return array
	 */
	public static function protected_resource_metadata(): array {
		return array(
			'resource'                 => self::resource_url(),
			'authorization_servers'    => array( self::issuer() ),
			'scopes_supported'         => self::scopes(),
			'bearer_methods_supported' => array( 'header' ),
			'resource_documentation'   => 'https://atlasaidev.com/docs/atlas-ai-connector/',
		);
	}

	/* --------------------------------------------------------------------- *
	 * 401 challenge
	 * --------------------------------------------------------------------- */

	/**
	 * Attach the OAuth challenge header to 401 responses from the MCP endpoint,
	 * so clients can discover the authorization server (MCP spec requirement).
	 *
	 * @param \WP_HTTP_Response $response Result to send.
	 * @param \WP_REST_Server   $server   Server instance.
	 * @param \WP_REST_Request  $request  Request used to generate the response.
	 * @return \WP_HTTP_Response
	 */
	public function add_challenge_header( $response, $server, $request ) {
		if ( ! self::is_enabled() || ! ( $response instanceof \WP_HTTP_Response ) ) {
			return $response;
		}
		if ( 401 !== (int) $response->get_status() ) {
			return $response;
		}
		$route = method_exists( $request, 'get_route' ) ? (string) $request->get_route() : '';
		if ( false === strpos( $route, '/mcp' ) ) {
			return $response;
		}
		$response->header(
			'WWW-Authenticate',
			sprintf( 'Bearer resource_metadata="%s", error="invalid_token"', esc_url_raw( self::protected_resource_metadata_url() ) )
		);
		return $response;
	}
}
