<?php

namespace AchlaAI;

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

final class Ownership {
	use Exact_Keys;

	/** @var Storage */
	private $storage;

	/** @var Api_Client_Interface */
	private $client;

	/** @var array<string,string> */
	private $public_keys;

	/** @var callable */
	private $clock;

	/**
	 * @param Storage                   $storage     Storage.
	 * @param Api_Client_Interface      $client      API client.
	 * @param array<string,string>|null $public_keys Pinned Ed25519 keys.
	 * @param callable|null             $clock       Unix timestamp provider.
	 */
	public function __construct( Storage $storage, Api_Client_Interface $client, $public_keys = null, $clock = null ) {
		$this->storage     = $storage;
		$this->client      = $client;
		$this->public_keys = null === $public_keys ? Trust_Anchors::ownership() : $public_keys;
		$this->clock       = null === $clock ? 'time' : $clock;
	}

	/** @return array<string,mixed>|\WP_Error */
	public function begin( $reconnect = false ) {
		$site = Protocol::canonical_site( home_url( '/' ) );
		if ( ! $site ) {
			return new \WP_Error( 'invalid_site_origin', __( 'This site does not have a supported public HTTP or HTTPS origin.', 'achlaai' ), array( 'status' => 400 ) );
		}
		$callback_uri = $this->callback_uri();
		$callback     = Protocol::canonical_site( $callback_uri );
		if ( ! $callback || $callback['site_key'] !== $site['site_key'] || wp_parse_url( $callback_uri, PHP_URL_PATH ) !== '/wp-json/achla/v1/ownership-callback' ) {
			return new \WP_Error( 'invalid_callback_uri', __( 'WordPress produced an unsupported ownership callback URL.', 'achlaai' ), array( 'status' => 400 ) );
		}

		$verifier = Crypto::random_token( 48 );
		$launch   = array(
			'protocol'              => Protocol::OWNERSHIP_PROTOCOL,
			'attempt_id'            => Crypto::random_token( 32 ),
			'state'                 => Crypto::random_token( 32 ),
			'code_challenge'        => Crypto::s256( $verifier ),
			'code_challenge_method' => Protocol::CHALLENGE_METHOD,
			'callback_uri'          => $callback_uri,
			'site'                  => $site['site_origin'],
			'finish_uri'            => Admin::settings_url(),
		);
		$now      = call_user_func( $this->clock );
		$attempt  = array(
			'status'             => 'pending',
			'created_at'         => $now,
			'expires_at'         => $now + Protocol::DNS_ATTEMPT_TTL,
			'verifier'           => $verifier,
			'launch'             => $launch,
			'launch_fingerprint' => Protocol::launch_fingerprint( $launch ),
			'last_error'         => '',
		);
		$this->save_attempt( $attempt );
		Lifecycle::schedule_dns_recovery( $now + Protocol::ATTEMPT_TTL );

		if ( $reconnect ) {
			$settings = $this->storage->get();
			if ( ! empty( $settings['connection']['binding_id'] ) ) {
				$settings['connection']['status'] = 'reverification_pending';
				$this->storage->save( $settings );
			}
		}

		return array(
			'ok'         => true,
			'status'     => 'pending',
			'launch_url' => $this->launch_url( $launch ),
		);
	}

	/** @return array<string,mixed>|\WP_Error */
	public function resume() {
		$attempt = $this->storage->get_attempt();
		$now     = call_user_func( $this->clock );
		if ( $attempt && 'complete' === $attempt['status'] ) {
			return array(
				'ok'     => true,
				'status' => 'complete',
			);
		}
		if ( $attempt && in_array( $attempt['status'], array( 'failed', 'expired' ), true ) ) {
			return $this->begin( true );
		}
		if ( $attempt && in_array( $attempt['status'], array( 'completing', 'callback_received' ), true )
			&& ! empty( $attempt['verifier'] ) && ! empty( $attempt['callback_snapshot'] ) && ! empty( $attempt['challenge_hash'] )
			&& (int) $attempt['expires_at'] > $now ) {
			return $this->complete( $attempt );
		}
		if ( ! $attempt || empty( $attempt['launch'] ) || empty( $attempt['verifier'] ) || (int) $attempt['expires_at'] <= $now ) {
			return $this->begin( true );
		}
		if ( ! Crypto::safe_equal( $attempt['launch_fingerprint'], Protocol::launch_fingerprint( $attempt['launch'] ) ) ) {
			$this->storage->delete_attempt();
			return new \WP_Error( 'mutated_launch', __( 'The saved setup attempt failed its integrity check. Start again.', 'achlaai' ), array( 'status' => 409 ) );
		}
		return array(
			'ok'         => true,
			'status'     => $attempt['status'],
			'launch_url' => $this->launch_url( $attempt['launch'] ),
		);
	}

	/** @return array<string,mixed>|\WP_Error */
	public function poll_dns_recovery( $force = false ) {
		$attempt = $this->storage->get_attempt();
		$now     = call_user_func( $this->clock );
		if ( ! $attempt || 'complete' === ( isset( $attempt['status'] ) ? $attempt['status'] : '' ) ) {
			return array(
				'ok'     => true,
				'status' => $attempt ? 'complete' : 'none',
			);
		}
		if ( (int) $attempt['expires_at'] <= $now ) {
			return $this->close_dns_attempt( $attempt, 'expired', 'dns_recovery_expired' );
		}
		if ( 'pending' === $attempt['status'] && $now < (int) $attempt['created_at'] + Protocol::ATTEMPT_TTL ) {
			Lifecycle::schedule_dns_recovery( (int) $attempt['created_at'] + Protocol::ATTEMPT_TTL );
			return array(
				'ok'     => true,
				'status' => 'pending',
			);
		}
		if ( ! in_array( $attempt['status'], array( 'pending', 'dns_pending' ), true ) ) {
			return array(
				'ok'     => true,
				'status' => $attempt['status'],
			);
		}
		if ( ! $this->valid_dns_attempt( $attempt ) ) {
			return $this->close_dns_attempt( $attempt, 'failed', 'mutated_launch' );
		}
		if ( ! $force && ! empty( $attempt['last_dns_poll_at'] ) && (int) $attempt['last_dns_poll_at'] + Protocol::DNS_POLL_INTERVAL > $now ) {
			Lifecycle::schedule_dns_recovery( (int) $attempt['last_dns_poll_at'] + Protocol::DNS_POLL_INTERVAL );
			return array(
				'ok'     => true,
				'status' => 'dns_pending',
			);
		}

		$attempt['status']           = 'dns_pending';
		$attempt['last_dns_poll_at'] = $now;
		$this->save_attempt( $attempt );
		$response = $this->client->complete_ownership_dns( $attempt['launch'], $attempt['verifier'] );
		if ( is_wp_error( $response ) ) {
			$data   = $response->get_error_data();
			$status = is_array( $data ) && isset( $data['status'] ) ? (int) $data['status'] : 0;
			if ( in_array( $status, array( 400, 409, 410 ), true ) ) {
				return $this->close_dns_attempt( $attempt, 'failed', $response->get_error_code() );
			}
			$attempt['last_error'] = $response->get_error_code();
			$this->save_attempt( $attempt );
			Lifecycle::schedule_dns_recovery( $now + Protocol::DNS_POLL_INTERVAL );
			return $response;
		}

		$site     = Protocol::canonical_site( $attempt['launch']['site'] );
		$snapshot = array(
			'project_id'     => isset( $response['project_id'] ) ? $response['project_id'] : '',
			'site_origin'    => $site ? $site['site_origin'] : '',
			'security_level' => $site ? $site['security_level'] : '',
		);
		$binding  = $this->validate_binding_response( $response, $snapshot );
		if ( is_wp_error( $binding ) ) {
			return $this->close_dns_attempt( $attempt, 'failed', $binding->get_error_code() );
		}
		$this->storage->store_binding( $binding );
		$this->storage->save_attempt(
			array(
				'status'     => 'complete',
				'expires_at' => $now + Protocol::ATTEMPT_TTL,
				'last_error' => '',
			),
			Protocol::ATTEMPT_TTL
		);
		wp_clear_scheduled_hook( Lifecycle::DNS_RECOVERY_CRON_HOOK );
		wp_schedule_single_event( $now + 5, Lifecycle::CRON_HOOK );
		wp_schedule_single_event( $now + 10, Lifecycle::HEARTBEAT_CRON_HOOK );
		return array(
			'ok'     => true,
			'status' => 'complete',
		);
	}

	/**
	 * @param string               $body_bytes Exact request bytes.
	 * @param array<string,string> $headers    Lowercase callback headers.
	 * @return array<string,mixed>|\WP_Error
	 */
	public function handle_callback( $body_bytes, $headers ) {
		if ( ! is_string( $body_bytes ) || '' === $body_bytes || strlen( $body_bytes ) > Protocol::MAX_CALLBACK_BYTES ) {
			return $this->callback_error( 'invalid_callback' );
		}
		$key_id    = isset( $headers['x-aip-key-id'] ) ? $headers['x-aip-key-id'] : '';
		$timestamp = isset( $headers['x-aip-timestamp'] ) ? $headers['x-aip-timestamp'] : '';
		$signature = isset( $headers['x-aip-signature'] ) ? $headers['x-aip-signature'] : '';
		if ( ! Protocol::is_identifier( $key_id ) || ! isset( $this->public_keys[ $key_id ] ) || 0 !== strpos( $signature, 'ed25519:' ) ) {
			return $this->callback_error( 'invalid_signature', 401 );
		}
		$payload = Protocol::callback_signature_payload( $this->callback_uri(), $body_bytes, $timestamp, $key_id );
		if ( ! Crypto::verify_ed25519( $payload, substr( $signature, 8 ), $this->public_keys[ $key_id ] ) ) {
			return $this->callback_error( 'invalid_signature', 401 );
		}

		$body = json_decode( $body_bytes, true );
		if ( ! $this->valid_callback_shape( $body, $timestamp ) ) {
			return $this->callback_error( 'invalid_callback' );
		}
		$attempt = $this->storage->get_attempt();
		if ( ! $attempt ) {
			return $this->callback_error( 'invalid_callback' );
		}

		$body_hash      = Crypto::sha256_base64url( $body_bytes );
		$challenge_hash = Crypto::sha256_base64url( $body['challenge'] );
		if ( 'complete' === $attempt['status'] ) {
			if ( Crypto::safe_equal( $attempt['callback_body_hash'], $body_hash ) && Crypto::safe_equal( $attempt['challenge_hash'], $challenge_hash ) ) {
				return array(
					'ok'       => true,
					'status'   => 'complete',
					'replayed' => true,
				);
			}
			return $this->callback_error( 'invalid_callback' );
		}

		if ( ! $this->callback_matches_attempt( $body, $attempt ) ) {
			return $this->callback_error( 'invalid_callback' );
		}
		$attempt['status']             = 'completing';
		$attempt['callback_challenge'] = $body['challenge'];
		$attempt['challenge_hash']     = $challenge_hash;
		$attempt['callback_body_hash'] = $body_hash;
		$attempt['callback_snapshot']  = $this->callback_snapshot( $body );
		$this->save_attempt( $attempt );

		return $this->complete( $attempt );
	}

	/** @return array<string,mixed> */
	public function status() {
		$settings       = $this->storage->get();
		$status         = $this->storage->reconcile_origin();
		$attempt        = $this->storage->get_attempt();
		$now            = call_user_func( $this->clock );
		$attempt_status = 'none';
		$current_site   = Protocol::canonical_site( home_url( '/' ) );
		if ( $attempt ) {
			$attempt_status = (int) $attempt['expires_at'] <= $now && 'complete' !== $attempt['status'] ? 'expired' : $attempt['status'];
		}
		return array(
			'connection'    => $status,
			'attempt'       => $attempt_status,
			'siteOrigin'    => $settings['connection']['site_origin'],
			'currentOrigin' => $current_site ? $current_site['site_origin'] : '',
			'securityLevel' => $settings['connection']['security_level'],
			'projectId'     => $settings['connection']['project_id'],
			'completedAt'   => $settings['connection']['completed_at'],
			'lastError'     => $attempt && ! empty( $attempt['last_error'] ) ? $attempt['last_error'] : '',
		);
	}

	/** @return true|\WP_Error */
	public function disconnect() {
		$settings = $this->storage->get();
		if ( empty( $settings['connection']['binding_id'] ) ) {
			return true;
		}
		$result = $this->client->disconnect( $settings['connection'] );
		if ( is_wp_error( $result ) ) {
			if ( $this->is_authoritative_rejection( $result ) ) {
				// The backend authoritatively rejected this binding (it no longer
				// recognizes it as owned by this installation). A transient network
				// error must preserve local state so the site keeps retrying, but an
				// authoritative 401/403 must not leave a stale token/key/binding/release
				// authority usable locally -- mirror Management_Diagnostics::report().
				$this->storage->revoke_management( 'disconnect_rejected' );
				$this->storage->delete_attempt();
			}
			return $result;
		}
		$defaults                           = $this->storage->defaults();
		$settings['connection']             = $defaults['connection'];
		$settings['lifecycle']              = $defaults['lifecycle'];
		$settings['release']                = $defaults['release'];
		$settings['management_diagnostics'] = $defaults['management_diagnostics'];
		$this->storage->save( $settings );
		$this->storage->delete_attempt();
		return true;
	}

	/** @return bool */
	private function is_authoritative_rejection( \WP_Error $error ) {
		$data = $error->get_error_data();
		return is_array( $data ) && isset( $data['status'] ) && in_array( (int) $data['status'], array( 401, 403 ), true );
	}

	/**
	 * @param array<string,mixed> $attempt Attempt.
	 * @return array<string,mixed>|\WP_Error
	 */
	private function complete( $attempt ) {
		$snapshot = $attempt['callback_snapshot'];
		$message  = Protocol::completion_message(
			array(
				'challenge_id'   => $snapshot['challenge_id'],
				'challenge_hash' => $attempt['challenge_hash'],
				'project_id'     => $snapshot['project_id'],
				'site_origin'    => $snapshot['site_origin'],
				'callback_uri'   => $snapshot['callback_uri'],
				'code_challenge' => $snapshot['code_challenge'],
			)
		);
		$proof    = Crypto::completion_proof( $attempt['verifier'], $message );
		$response = $this->client->complete_ownership( $snapshot['challenge_id'], $attempt['verifier'], $proof );
		if ( is_wp_error( $response ) ) {
			$attempt['status']     = 'callback_received';
			$attempt['last_error'] = $response->get_error_code();
			$this->save_attempt( $attempt );
			return $response;
		}
		$binding = $this->validate_binding_response( $response, $snapshot );
		if ( is_wp_error( $binding ) ) {
			$attempt['status']     = 'failed';
			$attempt['last_error'] = $binding->get_error_code();
			unset( $attempt['callback_challenge'] );
			$this->save_attempt( $attempt );
			return $binding;
		}

		$this->storage->store_binding( $binding );
		$complete_marker = array(
			'status'             => 'complete',
			'expires_at'         => call_user_func( $this->clock ) + Protocol::ATTEMPT_TTL,
			'challenge_hash'     => $attempt['challenge_hash'],
			'callback_body_hash' => $attempt['callback_body_hash'],
			'last_error'         => '',
		);
		$this->storage->save_attempt( $complete_marker );
		wp_clear_scheduled_hook( Lifecycle::DNS_RECOVERY_CRON_HOOK );
		return array(
			'ok'     => true,
			'status' => 'complete',
		);
	}

	/**
	 * @param mixed  $body      Parsed callback.
	 * @param string $timestamp Header timestamp.
	 * @return bool
	 */
	private function valid_callback_shape( $body, $timestamp ) {
		if ( ! is_array( $body ) || ! $this->has_exact_keys( $body, array( 'protocol', 'challenge_id', 'attempt_id', 'challenge', 'project_id', 'site_origin', 'callback_uri', 'code_challenge', 'issued_at', 'expires_at', 'aud', 'security_level' ) ) ) {
			return false;
		}
		$issued  = Protocol::parse_canonical_time( $body['issued_at'] );
		$expires = Protocol::parse_canonical_time( $body['expires_at'] );
		$now     = call_user_func( $this->clock );
		return Protocol::OWNERSHIP_PROTOCOL === $body['protocol']
			&& Protocol::OWNERSHIP_AUDIENCE === $body['aud']
			&& Protocol::is_token( $body['challenge_id'], 32, 128 )
			&& Protocol::is_token( $body['attempt_id'], 16, 256 )
			&& Protocol::is_token( $body['challenge'], 32, 256 )
			&& Protocol::is_identifier( $body['project_id'] )
			&& Protocol::is_token( $body['code_challenge'], 43, 43 )
			&& in_array( $body['security_level'], array( 'https', 'limited_http' ), true )
			&& Crypto::safe_equal( $timestamp, $body['issued_at'] )
			&& null !== $issued && null !== $expires && $expires > $issued
			&& $expires - $issued <= Protocol::CALLBACK_TTL
			&& $issued <= $now + Protocol::CALLBACK_TTL && $expires >= $now;
	}

	/**
	 * @param array<string,mixed> $body    Callback.
	 * @param array<string,mixed> $attempt Attempt.
	 * @return bool
	 */
	private function callback_matches_attempt( $body, $attempt ) {
		if ( empty( $attempt['launch'] ) || empty( $attempt['verifier'] ) || (int) $attempt['expires_at'] < call_user_func( $this->clock ) ) {
			return false;
		}
		$launch = $attempt['launch'];
		if ( ! Crypto::safe_equal( $attempt['launch_fingerprint'], Protocol::launch_fingerprint( $launch ) ) ) {
			return false;
		}
		$site = Protocol::canonical_site( $launch['site'] );
		return $site
			&& Crypto::safe_equal( $body['attempt_id'], $launch['attempt_id'] )
			&& Crypto::safe_equal( $body['site_origin'], $launch['site'] )
			&& Crypto::safe_equal( $body['callback_uri'], $launch['callback_uri'] )
			&& Crypto::safe_equal( $body['code_challenge'], $launch['code_challenge'] )
			&& Crypto::safe_equal( Crypto::s256( $attempt['verifier'] ), $launch['code_challenge'] )
			&& $body['security_level'] === $site['security_level'];
	}

	/** @return array<string,string> */
	private function callback_snapshot( $body ) {
		return array(
			'challenge_id'   => $body['challenge_id'],
			'project_id'     => $body['project_id'],
			'site_origin'    => $body['site_origin'],
			'callback_uri'   => $body['callback_uri'],
			'code_challenge' => $body['code_challenge'],
			'security_level' => $body['security_level'],
		);
	}

	/**
	 * @param array<string,mixed>  $response Completion response.
	 * @param array<string,string> $snapshot Callback snapshot.
	 * @return array<string,mixed>|\WP_Error
	 */
	private function validate_binding_response( $response, $snapshot ) {
		if ( ! $this->has_exact_keys( $response, array( 'protocol', 'binding_id', 'installation_id', 'project_id', 'widget', 'site_origin', 'security_level', 'completed_at', 'management' ) )
			|| ! isset( $response['widget'] ) || ! is_array( $response['widget'] ) || ! $this->has_exact_keys( $response['widget'], array( 'pk', 'mode' ) )
			|| ! isset( $response['management'] ) || ! is_array( $response['management'] )
			|| ! $this->has_exact_keys( $response['management'], array( 'protocol', 'token_type', 'token', 'heartbeat_path', 'disconnect_path', 'diagnostics_path', 'heartbeat_interval_seconds', 'stale_after_days' ) ) ) {
			return new \WP_Error( 'invalid_binding_response', __( 'Achla returned an invalid binding response.', 'achlaai' ), array( 'status' => 502 ) );
		}
		$site       = Protocol::canonical_site( $response['site_origin'] );
		$management = $response['management'];
		if ( Protocol::OWNERSHIP_PROTOCOL !== $response['protocol']
			|| ! Protocol::is_identifier( $response['binding_id'] )
			|| ! Protocol::is_identifier( $response['installation_id'] )
			|| ! Protocol::is_identifier( $response['project_id'] )
			|| ! Crypto::safe_equal( $response['project_id'], $snapshot['project_id'] )
			|| ! is_string( $response['widget']['pk'] ) || ! preg_match( '/^pk_[A-Za-z0-9_-]+$/', $response['widget']['pk'] )
			|| 'bubble' !== $response['widget']['mode']
			|| ! $site || ! Crypto::safe_equal( $site['site_origin'], $snapshot['site_origin'] )
			|| $response['security_level'] !== $snapshot['security_level']
			|| 'aip-management-v1' !== $management['protocol']
			|| 'Bearer' !== $management['token_type']
			|| ! is_string( $management['token'] ) || 1 !== preg_match( '/^aipm_[A-Za-z0-9_-]{43}$/', $management['token'] )
			|| '/api/plugin/heartbeat' !== $management['heartbeat_path']
			|| '/api/plugin/disconnect' !== $management['disconnect_path']
			|| '/api/plugin/diagnostics' !== $management['diagnostics_path']
			|| 3600 !== $management['heartbeat_interval_seconds']
			|| 30 !== $management['stale_after_days']
			|| null === Protocol::parse_canonical_time( $response['completed_at'] ) ) {
			return new \WP_Error( 'invalid_binding_response', __( 'Achla returned an invalid binding response.', 'achlaai' ), array( 'status' => 502 ) );
		}
		return array(
			'binding_id'      => $response['binding_id'],
			'installation_id' => $response['installation_id'],
			'project_id'      => $response['project_id'],
			'widget_pk'       => $response['widget']['pk'],
			'site_origin'     => $site['site_origin'],
			'site_key'        => $site['site_key'],
			'security_level'  => $response['security_level'],
			'completed_at'    => $response['completed_at'],
			'management'      => $management,
		);
	}

	/** @return bool */
	private function valid_dns_attempt( $attempt ) {
		if ( empty( $attempt['launch'] ) || ! is_array( $attempt['launch'] ) || empty( $attempt['launch_fingerprint'] ) || empty( $attempt['verifier'] )
			|| ! $this->has_exact_keys( $attempt['launch'], array( 'protocol', 'attempt_id', 'state', 'code_challenge', 'code_challenge_method', 'callback_uri', 'site', 'finish_uri' ) )
			|| ! Protocol::is_token( $attempt['verifier'], 43, 128 ) ) {
			return false;
		}
		return Crypto::safe_equal( $attempt['launch_fingerprint'], Protocol::launch_fingerprint( $attempt['launch'] ) )
			&& Crypto::safe_equal( Crypto::s256( $attempt['verifier'] ), $attempt['launch']['code_challenge'] );
	}

	/** @return array<string,mixed> */
	private function close_dns_attempt( $attempt, $status, $error ) {
		$attempt['status']     = $status;
		$attempt['last_error'] = $error;
		unset( $attempt['verifier'] );
		$this->save_attempt( $attempt );
		wp_clear_scheduled_hook( Lifecycle::DNS_RECOVERY_CRON_HOOK );
		return array(
			'ok'     => false,
			'status' => $status,
			'error'  => $error,
		);
	}

	private function save_attempt( $attempt ) {
		$remaining = max( 1, (int) $attempt['expires_at'] - (int) call_user_func( $this->clock ) );
		$this->storage->save_attempt( $attempt, $remaining );
	}

	/** @return string */
	private function callback_uri() {
		return rest_url( 'achla/v1/ownership-callback' );
	}

	/**
	 * @param array<string,string> $launch Launch.
	 * @return string
	 */
	private function launch_url( $launch ) {
		return add_query_arg( $launch, ACHLAAI_APP_URL . '/en/connect' );
	}

	/** @return \WP_Error */
	private function callback_error( $code, $status = 400 ) {
		return new \WP_Error( $code, 'Ownership callback rejected.', array( 'status' => $status ) );
	}
}