<?php
/**
 * Amazon Bedrock credential resolution.
 *
 * Resolution order:
 *   1. wp-config.php constants
 *   2. Encrypted plugin options
 *   3. Process environment variables
 *   4. ECS/EKS container credential endpoint
 *   5. EC2 instance metadata service (IMDSv2)
 *
 * Sources 3-5 remove the need to store long-lived AWS keys in WordPress.
 *
 * An Amazon Bedrock API key is resolved separately, because it is not a signing credential:
 * it is sent as a bearer token and only Bedrock and Bedrock Runtime accept it. See api_key().
 *
 * @package AI_Chat_Bedrock
 */

if ( ! defined( 'ABSPATH' ) ) {
	exit;
}

class AI_Chat_Bedrock_AWS_Credentials {

	const IMDS_HOST          = '169.254.169.254';
	const CONTAINER_HOST     = '169.254.170.2';
	const ROLE_CACHE_KEY     = 'aicfab_role_credentials';
	const ROLE_CACHE_MIN_TTL = 60;
	const ROLE_CACHE_MAX_TTL = 900;
	const ROLE_EXPIRY_BUFFER = 300;

	/*
	 * After no IAM role answered, how long before asking again. Off AWS, or on an instance
	 * without a role, the metadata service never answers, and every request waited for it to
	 * time out until credentials were entered: a front page with the chat, an admin screen,
	 * a cron run. Saving the settings or running Diagnostics asks again at once.
	 */
	const ROLE_MISS_KEY = 'aicfab_role_credentials_miss';
	const ROLE_MISS_TTL = 300;

	/**
	 * Whether no role answered earlier in this request, for sites without a persistent cache
	 * where the transient is still read from the database each time.
	 *
	 * @var bool
	 */
	private static $role_missed = false;

	/**
	 * Resolve credentials for signing Bedrock requests.
	 *
	 * @param array $options Plugin options.
	 * @return array|WP_Error Credential array with access_key, secret_key, session_token and source.
	 */
	public static function resolve( $options = null ) {
		if ( ! is_array( $options ) ) {
			$options = get_option( 'ai_chat_bedrock_settings', array() );
			$options = is_array( $options ) ? $options : array();
		}

		$static = self::from_constants();
		if ( null !== $static ) {
			return $static;
		}

		$static = self::from_options( $options );
		if ( null !== $static ) {
			return $static;
		}

		if ( ! self::role_credentials_enabled( $options ) ) {
			return new WP_Error( 'aicfab_no_credentials', __( 'Amazon Bedrock credentials are not configured.', 'ai-chat-for-amazon-bedrock' ) );
		}

		$cached = self::cached_role_credentials();
		if ( null !== $cached ) {
			return $cached;
		}

		$credentials = self::from_environment();
		if ( null !== $credentials ) {
			return $credentials;
		}

		if ( ! self::$role_missed && false !== get_transient( self::ROLE_MISS_KEY ) ) {
			self::$role_missed = true;
		}
		if ( ! self::$role_missed ) {
			foreach ( array( 'from_container_role', 'from_instance_role' ) as $provider ) {
				$credentials = self::{$provider}();
				if ( null !== $credentials ) {
					self::cache_role_credentials( $credentials );
					return $credentials;
				}
			}
			self::$role_missed = true;
			set_transient( self::ROLE_MISS_KEY, 1, self::ROLE_MISS_TTL );
		}

		return new WP_Error( 'aicfab_no_credentials', __( 'Amazon Bedrock credentials are not configured.', 'ai-chat-for-amazon-bedrock' ) );
	}

	/**
	 * Resolve the Amazon Bedrock API key to use, if any.
	 *
	 * An API key is the quickest way to connect: one value, created in the Bedrock console,
	 * with no IAM user or access key pair. Resolution order is the AI_CHAT_BEDROCK_API_KEY
	 * constant, then the encrypted setting, then the key entered on Settings → Connectors
	 * (WordPress 7.1 and later, also encrypted), then the AWS_BEARER_TOKEN_BEDROCK
	 * environment variable the AWS SDKs read. The environment variable is ignored when access keys are
	 * configured in wp-config.php or the settings, because a value the site owner entered for
	 * this plugin should not be overridden by one the server happens to carry.
	 *
	 * @param array $options Plugin options.
	 * @return array|null Array with key, source and temporary, or null.
	 */
	public static function api_key( $options = null ) {
		if ( ! is_array( $options ) ) {
			$options = get_option( 'ai_chat_bedrock_settings', array() );
			$options = is_array( $options ) ? $options : array();
		}

		$candidates = array();
		if ( defined( 'AI_CHAT_BEDROCK_API_KEY' ) ) {
			$candidates['api_key_constant'] = (string) constant( 'AI_CHAT_BEDROCK_API_KEY' );
		}
		if ( ! empty( $options['bedrock_api_key'] ) ) {
			$candidates['api_key_option'] = AI_Chat_Bedrock_Security::decrypt_secret( $options['bedrock_api_key'] );
		}
		if ( class_exists( 'AI_Chat_Bedrock_Core_AI' ) ) {
			$candidates['api_key_connector'] = AI_Chat_Bedrock_Core_AI::connector_api_key();
		}
		$candidates['api_key_environment'] = self::env( 'AWS_BEARER_TOKEN_BEDROCK' );

		foreach ( $candidates as $source => $value ) {
			$key = self::clean_api_key( $value );
			if ( '' === $key ) {
				continue;
			}
			if ( 'api_key_environment' === $source && self::has_static_keys( $options ) ) {
				return null;
			}
			return array(
				'key'       => $key,
				'source'    => $source,
				'temporary' => 0 === strpos( $key, 'bedrock-api-key-' ),
			);
		}
		return null;
	}

	/**
	 * Normalise a pasted Amazon Bedrock API key.
	 *
	 * Long-term keys are base64 text and short-term keys are a prefixed, encoded presigned
	 * URL, so both fit a token character set. A pasted "Bearer " prefix or surrounding
	 * whitespace is removed. Anything else is rejected rather than sent in a header.
	 *
	 * @param string $value Raw value.
	 * @return string Clean key, or an empty string.
	 */
	public static function clean_api_key( $value ) {
		$value = trim( (string) $value );
		$value = preg_replace( '/^Bearer\s+/i', '', $value );
		if ( strlen( $value ) < 20 || strlen( $value ) > 8192 ) {
			return '';
		}
		return 1 === preg_match( '/^[A-Za-z0-9+\/=._~:-]+$/', $value ) ? $value : '';
	}

	/**
	 * Whether AWS access keys are configured in wp-config.php or the settings.
	 *
	 * @param array $options Plugin options.
	 * @return bool
	 */
	public static function has_static_keys( $options ) {
		return null !== self::from_constants() || null !== self::from_options( is_array( $options ) ? $options : array() );
	}

	/**
	 * Describe the active credential source without exposing secret values.
	 *
	 * @param array $options Plugin options.
	 * @return array
	 */
	public static function describe( $options = null ) {
		$api_key = self::api_key( $options );
		if ( null !== $api_key ) {
			return array(
				'configured' => true,
				'source'     => $api_key['source'],
				'temporary'  => $api_key['temporary'],
				'message'    => self::source_label( $api_key['source'] ),
			);
		}

		$credentials = self::resolve( $options );
		if ( is_wp_error( $credentials ) ) {
			return array(
				'configured' => false,
				'source'     => 'none',
				'temporary'  => false,
				'message'    => $credentials->get_error_message(),
			);
		}

		return array(
			'configured' => true,
			'source'     => $credentials['source'],
			'temporary'  => '' !== $credentials['session_token'],
			'message'    => self::source_label( $credentials['source'] ),
		);
	}

	/**
	 * Whether role-based credential providers may be used.
	 *
	 * @param array $options Plugin options.
	 * @return bool
	 */
	public static function role_credentials_enabled( $options ) {
		$enabled = ! isset( $options['aws_use_role_credentials'] ) || ! empty( $options['aws_use_role_credentials'] );
		return (bool) apply_filters( 'ai_chat_bedrock_use_role_credentials', $enabled );
	}

	/**
	 * Human readable label for a credential source.
	 *
	 * @param string $source Source identifier.
	 * @return string
	 */
	public static function source_label( $source ) {
		$labels = array(
			'constants'           => __( 'wp-config.php constants', 'ai-chat-for-amazon-bedrock' ),
			'options'             => __( 'Encrypted WordPress settings', 'ai-chat-for-amazon-bedrock' ),
			'environment'         => __( 'Server environment variables', 'ai-chat-for-amazon-bedrock' ),
			'container_role'      => __( 'ECS or EKS task role', 'ai-chat-for-amazon-bedrock' ),
			'instance_role'       => __( 'EC2 instance role (IMDSv2)', 'ai-chat-for-amazon-bedrock' ),
			'api_key_constant'    => __( 'Amazon Bedrock API key (wp-config.php constant)', 'ai-chat-for-amazon-bedrock' ),
			'api_key_option'      => __( 'Amazon Bedrock API key (encrypted WordPress settings)', 'ai-chat-for-amazon-bedrock' ),
			'api_key_connector'   => __( 'Amazon Bedrock API key (Settings > Connectors, encrypted)', 'ai-chat-for-amazon-bedrock' ),
			'api_key_environment' => __( 'Amazon Bedrock API key (AWS_BEARER_TOKEN_BEDROCK)', 'ai-chat-for-amazon-bedrock' ),
		);
		return isset( $labels[ $source ] ) ? $labels[ $source ] : __( 'Unknown source', 'ai-chat-for-amazon-bedrock' );
	}

	/**
	 * Discard cached role credentials, and the record that no role answered, so the next
	 * request asks the metadata service again.
	 */
	public static function flush_cache() {
		delete_transient( self::ROLE_CACHE_KEY );
		delete_transient( self::ROLE_MISS_KEY );
		self::$role_missed = false;
	}

	private static function from_constants() {
		if ( ! defined( 'AI_CHAT_BEDROCK_AWS_ACCESS_KEY' ) || ! defined( 'AI_CHAT_BEDROCK_AWS_SECRET_KEY' ) ) {
			return null;
		}
		$access = trim( (string) constant( 'AI_CHAT_BEDROCK_AWS_ACCESS_KEY' ) );
		$secret = trim( (string) constant( 'AI_CHAT_BEDROCK_AWS_SECRET_KEY' ) );
		$token  = defined( 'AI_CHAT_BEDROCK_AWS_SESSION_TOKEN' ) ? trim( (string) constant( 'AI_CHAT_BEDROCK_AWS_SESSION_TOKEN' ) ) : '';
		return self::build( $access, $secret, $token, 'constants' );
	}

	private static function from_options( $options ) {
		$access = isset( $options['aws_access_key'] ) ? AI_Chat_Bedrock_Security::decrypt_secret( $options['aws_access_key'] ) : '';
		$secret = isset( $options['aws_secret_key'] ) ? AI_Chat_Bedrock_Security::decrypt_secret( $options['aws_secret_key'] ) : '';
		$token  = isset( $options['aws_session_token'] ) ? AI_Chat_Bedrock_Security::decrypt_secret( $options['aws_session_token'] ) : '';
		return self::build( $access, $secret, $token, 'options' );
	}

	private static function from_environment() {
		$access = self::env( 'AWS_ACCESS_KEY_ID' );
		$secret = self::env( 'AWS_SECRET_ACCESS_KEY' );
		$token  = self::env( 'AWS_SESSION_TOKEN' );
		return self::build( $access, $secret, $token, 'environment' );
	}

	private static function from_container_role() {
		$relative = self::env( 'AWS_CONTAINER_CREDENTIALS_RELATIVE_URI' );
		$full     = self::env( 'AWS_CONTAINER_CREDENTIALS_FULL_URI' );

		if ( '' !== $relative ) {
			$url = 'http://' . self::CONTAINER_HOST . '/' . ltrim( $relative, '/' );
		} elseif ( '' !== $full && self::is_allowed_container_uri( $full ) ) {
			$url = $full;
		} else {
			return null;
		}

		$headers = array( 'Accept' => 'application/json' );
		$token   = self::container_authorization_token();
		if ( '' !== $token ) {
			$headers['Authorization'] = $token;
		}

		$body = self::request( 'GET', $url, $headers );
		return null === $body ? null : self::from_metadata_document( $body, 'container_role' );
	}

	private static function from_instance_role() {
		$base  = 'http://' . self::IMDS_HOST . '/latest';
		$token = self::request(
			'PUT',
			$base . '/api/token',
			array( 'X-aws-ec2-metadata-token-ttl-seconds' => '60' )
		);
		if ( null === $token || '' === trim( $token ) ) {
			return null;
		}
		$headers = array( 'X-aws-ec2-metadata-token' => trim( $token ) );

		$role = self::request( 'GET', $base . '/meta-data/iam/security-credentials/', $headers );
		if ( null === $role ) {
			return null;
		}
		$role = trim( strtok( trim( $role ), "\n" ) );
		if ( '' === $role || ! preg_match( '#^[A-Za-z0-9+=,.@_/-]{1,128}$#', $role ) ) {
			return null;
		}

		$document = self::request( 'GET', $base . '/meta-data/iam/security-credentials/' . rawurlencode( $role ), $headers );
		return null === $document ? null : self::from_metadata_document( $document, 'instance_role' );
	}

	private static function from_metadata_document( $body, $source ) {
		$data = json_decode( (string) $body, true );
		if ( ! is_array( $data ) ) {
			return null;
		}
		if ( isset( $data['Code'] ) && 'Success' !== $data['Code'] ) {
			return null;
		}

		$credentials = self::build(
			isset( $data['AccessKeyId'] ) ? (string) $data['AccessKeyId'] : '',
			isset( $data['SecretAccessKey'] ) ? (string) $data['SecretAccessKey'] : '',
			isset( $data['Token'] ) ? (string) $data['Token'] : '',
			$source
		);
		if ( null === $credentials || '' === $credentials['session_token'] ) {
			return null;
		}

		if ( ! empty( $data['Expiration'] ) ) {
			$expires = strtotime( (string) $data['Expiration'] );
			if ( false !== $expires ) {
				$credentials['expires'] = $expires;
			}
		}
		return $credentials;
	}

	private static function cached_role_credentials() {
		$cached = get_transient( self::ROLE_CACHE_KEY );
		if ( ! is_array( $cached ) || empty( $cached['access_key'] ) || empty( $cached['secret_key'] ) ) {
			return null;
		}

		$access = AI_Chat_Bedrock_Security::decrypt_secret( $cached['access_key'] );
		$secret = AI_Chat_Bedrock_Security::decrypt_secret( $cached['secret_key'] );
		$token  = isset( $cached['session_token'] ) ? AI_Chat_Bedrock_Security::decrypt_secret( $cached['session_token'] ) : '';
		$source = isset( $cached['source'] ) ? sanitize_key( $cached['source'] ) : 'instance_role';

		if ( isset( $cached['expires'] ) && (int) $cached['expires'] - self::ROLE_EXPIRY_BUFFER <= time() ) {
			self::flush_cache();
			return null;
		}
		return self::build( $access, $secret, $token, $source );
	}

	private static function cache_role_credentials( $credentials ) {
		$access = AI_Chat_Bedrock_Security::encrypt_secret( $credentials['access_key'] );
		$secret = AI_Chat_Bedrock_Security::encrypt_secret( $credentials['secret_key'] );
		$token  = AI_Chat_Bedrock_Security::encrypt_secret( $credentials['session_token'] );
		if ( '' === $access || '' === $secret || '' === $token ) {
			return;
		}

		$ttl = self::ROLE_CACHE_MAX_TTL;
		if ( isset( $credentials['expires'] ) ) {
			$ttl = (int) $credentials['expires'] - self::ROLE_EXPIRY_BUFFER - time();
		}
		$ttl = max( self::ROLE_CACHE_MIN_TTL, min( self::ROLE_CACHE_MAX_TTL, $ttl ) );

		$payload = array(
			'access_key'    => $access,
			'secret_key'    => $secret,
			'session_token' => $token,
			'source'        => $credentials['source'],
		);
		if ( isset( $credentials['expires'] ) ) {
			$payload['expires'] = (int) $credentials['expires'];
		}
		set_transient( self::ROLE_CACHE_KEY, $payload, $ttl );
	}

	/**
	 * Perform a metadata request.
	 *
	 * The IMDS and container credential endpoints are link-local addresses that
	 * WordPress safe-HTTP helpers deliberately reject, so these fixed, non
	 * user-controlled URLs are requested directly with redirects disabled.
	 *
	 * @param string $method  HTTP method.
	 * @param string $url     Metadata URL.
	 * @param array  $headers Request headers.
	 * @return string|null
	 */
	private static function request( $method, $url, $headers = array() ) {
		$args = array(
			'method'             => $method,
			'timeout'            => 2,
			'redirection'        => 0,
			'httpversion'        => '1.1',
			'reject_unsafe_urls' => false,
			'sslverify'          => false,
			'headers'            => $headers,
		);

		$response = wp_remote_request( $url, $args );
		if ( is_wp_error( $response ) ) {
			return null;
		}
		$status = (int) wp_remote_retrieve_response_code( $response );
		if ( $status < 200 || $status >= 300 ) {
			return null;
		}
		$body = wp_remote_retrieve_body( $response );
		return strlen( $body ) > 8192 ? null : $body;
	}

	private static function container_authorization_token() {
		$file = self::env( 'AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE' );
		if ( '' !== $file && 0 === strpos( $file, '/' ) && is_readable( $file ) && filesize( $file ) <= 8192 ) {
			$contents = file_get_contents( $file ); // phpcs:ignore WordPress.WP.AlternativeFunctions.file_get_contents_file_get_contents
			if ( false !== $contents ) {
				return trim( $contents );
			}
		}
		return self::env( 'AWS_CONTAINER_AUTHORIZATION_TOKEN' );
	}

	private static function is_allowed_container_uri( $uri ) {
		$parts = wp_parse_url( $uri );
		if ( ! is_array( $parts ) || empty( $parts['host'] ) || empty( $parts['scheme'] ) ) {
			return false;
		}
		$scheme = strtolower( $parts['scheme'] );
		$host   = strtolower( $parts['host'] );
		if ( 'https' === $scheme ) {
			return true;
		}
		if ( 'http' !== $scheme ) {
			return false;
		}
		return in_array( $host, array( self::CONTAINER_HOST, 'localhost', '127.0.0.1', '[::1]', '::1' ), true );
	}

	private static function env( $name ) {
		$value = getenv( $name );
		if ( false === $value && isset( $_SERVER[ $name ] ) ) {
			// A server variable, not request input, but it is still unslashed and cleaned.
			$value = sanitize_text_field( wp_unslash( $_SERVER[ $name ] ) );
		}
		return false === $value ? '' : trim( (string) $value );
	}

	private static function build( $access_key, $secret_key, $session_token, $source ) {
		$access_key    = self::clean( $access_key );
		$secret_key    = self::clean( $secret_key );
		$session_token = self::clean( $session_token );

		if ( '' === $access_key || '' === $secret_key ) {
			return null;
		}
		return array(
			'access_key'    => $access_key,
			'secret_key'    => $secret_key,
			'session_token' => $session_token,
			'source'        => $source,
		);
	}

	private static function clean( $value ) {
		$value = trim( (string) $value );
		return preg_replace( '/[\r\n\t\0\x0B]/', '', $value );
	}
}