<?php
/**
 * Vector format migration (v1 → v2).
 *
 * v1: 3072-dim OpenAI vector, JSON string in `_ai_search_embedding` postmeta.
 * v2: 1024-dim L2-normalized float32, base64(pack('f*', …)) in `_ai_search_embedding_v2`.
 *
 * The migration is READ-ONLY on the v1 data:
 *   - Writes v2 to a NEW meta key
 *   - Leaves `_ai_search_embedding` (v1) untouched
 *   - Only flips `ai_search_vector_format_version` from '1' to '2' when every
 *     post has both keys populated
 *   - Rollback = flip the option back; the plugin reads v1 again
 *
 * Why 1024 dims:
 *   OpenAI text-embedding-3-* supports Matryoshka Representation Learning —
 *   the first N dims of the vector are themselves a meaningful embedding at
 *   that smaller dimension, provided we L2-renormalize after the slice.
 *   1024 keeps ~98% of MTEB quality while cutting cosine CPU + storage by 3×.
 *
 * Why packed binary + base64:
 *   - 1024 float32 = 4096 bytes raw
 *   - base64 adds ~33% → ~5460 chars per postmeta row
 *   - vs ~25 000 chars of JSON today = 4.5× smaller postmeta I/O
 *   - unpack('f*', …) is ~10× faster than json_decode on 1024 floats
 *
 * @since 1.45.0
 */

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

class AI_Search_Vector_Migration {

	const META_KEY_V1          = '_ai_search_embedding';
	const META_KEY_V2          = '_ai_search_embedding_v2';
	const TARGET_DIM           = 1024;
	const BATCH_SIZE           = 50;
	const OPTION_VERSION       = 'ai_search_vector_format_version';
	const OPTION_DISMISSED     = 'ai_search_v2_migration_dismissed_until';
	const OPTION_STARTED_AT    = 'ai_search_v2_migration_started_at';
	const OPTION_FINISHED_AT   = 'ai_search_v2_migration_finished_at';
	const TRANSIENT_TOTAL      = 'ai_search_v2_migration_total';
	const CRON_HOOK            = 'ai_search_v2_migration_cron';
	const CRON_LOCK_TRANSIENT  = 'ai_search_v2_migration_cron_lock';
	const CRON_BUDGET_SECONDS  = 10;
	const DONE_NOTICE_KEY      = 'ai_search_v2_auto_migrated_notice';

	/**
	 * Current stored format version. '1' or '2'.
	 */
	public static function current_version(): string {
		return (string) get_option( self::OPTION_VERSION, '1' );
	}

	/**
	 * True when the plugin should read + write the v2 format.
	 */
	public static function is_v2_active(): bool {
		return '2' === self::current_version();
	}

	/**
	 * True when we have at least one v1 embedding without a matching v2,
	 * meaning the admin should see the migration card.
	 */
	public static function needs_migration(): bool {
		if ( self::is_v2_active() ) {
			return false;
		}
		return self::count_pending() > 0;
	}

	/**
	 * Count posts with v1 but no v2 — i.e. still to migrate.
	 */
	public static function count_pending(): int {
		global $wpdb;
		$sql = $wpdb->prepare(
			"SELECT COUNT(DISTINCT pm1.post_id)
			 FROM {$wpdb->postmeta} pm1
			 LEFT JOIN {$wpdb->postmeta} pm2
			       ON pm2.post_id = pm1.post_id AND pm2.meta_key = %s
			 WHERE pm1.meta_key = %s
			   AND pm2.meta_id IS NULL",
			self::META_KEY_V2,
			self::META_KEY_V1
		);
		return (int) $wpdb->get_var( $sql );
	}

	/**
	 * Total v1 embeddings in the catalog (= migration target size).
	 */
	public static function count_total(): int {
		$cached = get_transient( self::TRANSIENT_TOTAL );
		if ( false !== $cached ) {
			return (int) $cached;
		}
		global $wpdb;
		$total = (int) $wpdb->get_var( $wpdb->prepare(
			"SELECT COUNT(*) FROM {$wpdb->postmeta} WHERE meta_key = %s",
			self::META_KEY_V1
		) );
		set_transient( self::TRANSIENT_TOTAL, $total, HOUR_IN_SECONDS );
		return $total;
	}

	/**
	 * Progress snapshot for UI.
	 */
	public static function progress(): array {
		$total   = self::count_total();
		$pending = self::count_pending();
		$done    = max( 0, $total - $pending );
		$pct     = $total > 0 ? min( 100, (int) round( $done * 100 / $total ) ) : 100;
		return [
			'total'   => $total,
			'done'    => $done,
			'pending' => $pending,
			'pct'     => $pct,
			'version' => self::current_version(),
		];
	}

	/* -------------------------------------------------------------
	 * Encoding helpers — exposed statically so the hot search path
	 * can call them without a round-trip through this class' state.
	 * ----------------------------------------------------------- */

	/**
	 * Matryoshka truncation + L2 renormalization.
	 *
	 * Keeps the first $dim components (which carry the most information
	 * by construction of the training objective) and renormalizes so the
	 * resulting sub-vector is unit length. This is what makes smaller
	 * dims a valid embedding at all.
	 *
	 * @param float[] $v
	 * @param int     $dim
	 * @return float[] array of length $dim, L2 norm = 1 (or [] if zero vector)
	 */
	public static function truncate_and_normalize( array $v, int $dim = self::TARGET_DIM ): array {
		if ( count( $v ) < $dim ) {
			$dim = count( $v );
		}
		$slice = array_slice( $v, 0, $dim );
		$sq    = 0.0;
		foreach ( $slice as $x ) {
			$sq += $x * $x;
		}
		// v1.45.0 fix: reject the zero vector explicitly. A 1024-element
		// array of 0.0 would otherwise slip through the empty() check in
		// migrate_batch() and be stored as a "valid" packed vector.
		// OpenAI never returns pure zeros, but defence in depth.
		if ( $sq <= 0 ) {
			return [];
		}
		$inv = 1.0 / sqrt( $sq );
		foreach ( $slice as &$x ) {
			$x = (float) ( $x * $inv );
		}
		return $slice;
	}

	/**
	 * Pack a 1024-dim float array as base64(pack('f*',…)) for postmeta storage.
	 *
	 * @param float[] $v
	 * @return string
	 */
	public static function pack_vector( array $v ): string {
		return base64_encode( pack( 'f*', ...$v ) );
	}

	/**
	 * Inverse of pack_vector. Returns [] on bad input.
	 *
	 * @param string $b64
	 * @return float[]
	 */
	public static function unpack_vector( string $b64 ): array {
		if ( '' === $b64 ) {
			return [];
		}
		$bin = base64_decode( $b64, true );
		if ( false === $bin || 0 !== strlen( $bin ) % 4 ) {
			return [];
		}
		$arr = unpack( 'f*', $bin );
		return $arr ? array_values( $arr ) : [];
	}

	/**
	 * Fast dot product between two equal-length float arrays.
	 *
	 * When both vectors are L2-normalized (v2 vectors always are), this
	 * is identical to cosine similarity and skips the two sqrt() calls
	 * the generic calculate_similarity() does.
	 *
	 * @return float
	 */
	public static function dot( array $a, array $b ): float {
		$sum = 0.0;
		$n   = min( count( $a ), count( $b ) );
		for ( $i = 0; $i < $n; $i++ ) {
			$sum += $a[ $i ] * $b[ $i ];
		}
		return (float) $sum;
	}

	/* -------------------------------------------------------------
	 * Migration batch — the loop that converts v1 → v2 on disk.
	 * ----------------------------------------------------------- */

	/**
	 * Convert up to BATCH_SIZE pending posts.
	 *
	 * Idempotent: skips posts that already have a v2 meta. Safe to re-run.
	 *
	 * @return array { processed, failed[], remaining, elapsed_ms }
	 */
	public static function migrate_batch(): array {
		global $wpdb;
		$t0 = microtime( true );

		$ids = $wpdb->get_col( $wpdb->prepare(
			"SELECT pm1.post_id
			 FROM {$wpdb->postmeta} pm1
			 LEFT JOIN {$wpdb->postmeta} pm2
			       ON pm2.post_id = pm1.post_id AND pm2.meta_key = %s
			 WHERE pm1.meta_key = %s
			   AND pm2.meta_id IS NULL
			 ORDER BY pm1.post_id ASC
			 LIMIT %d",
			self::META_KEY_V2,
			self::META_KEY_V1,
			self::BATCH_SIZE
		) );

		$processed = 0;
		$failed    = [];

		foreach ( $ids as $pid ) {
			$pid     = (int) $pid;
			$raw     = get_post_meta( $pid, self::META_KEY_V1, true );
			$decoded = json_decode( (string) $raw, true );

			// Older rows wrapped the vector as `{ embedding: [...] }`.
			if ( is_array( $decoded ) && isset( $decoded['embedding'] ) && is_array( $decoded['embedding'] ) ) {
				$decoded = $decoded['embedding'];
			}

			if ( ! is_array( $decoded ) || count( $decoded ) < 2 ) {
				$failed[] = [ 'id' => $pid, 'reason' => 'v1_malformed' ];
				continue;
			}

			$truncated = self::truncate_and_normalize( $decoded, self::TARGET_DIM );
			if ( empty( $truncated ) ) {
				$failed[] = [ 'id' => $pid, 'reason' => 'zero_vector' ];
				continue;
			}

			$packed = self::pack_vector( $truncated );
			update_post_meta( $pid, self::META_KEY_V2, $packed );
			$processed++;
		}

		// Bust the cached total count so progress UI reflects reality.
		delete_transient( self::TRANSIENT_TOTAL );

		return [
			'processed'  => $processed,
			'failed'     => $failed,
			'remaining'  => self::count_pending(),
			'elapsed_ms' => (int) round( ( microtime( true ) - $t0 ) * 1000 ),
		];
	}

	/**
	 * Mark the migration as complete and switch the plugin to v2 reads.
	 * No-op if there are still pending rows (safety).
	 */
	public static function finalize(): bool {
		if ( self::count_pending() > 0 ) {
			return false;
		}
		update_option( self::OPTION_VERSION, '2', false );
		update_option( self::OPTION_FINISHED_AT, time(), false );
		delete_transient( self::TRANSIENT_TOTAL );
		return true;
	}

	/**
	 * Rollback: flip version back to '1'. The v2 postmeta stays on disk
	 * (harmless — ignored while version=1) so a second migration is instant.
	 */
	public static function rollback(): void {
		update_option( self::OPTION_VERSION, '1', false );
		delete_option( self::OPTION_FINISHED_AT );
	}

	/* -------------------------------------------------------------
	 * AJAX handlers — registered by AI_Search_Admin_Manager.
	 * ----------------------------------------------------------- */

	public static function ajax_start(): void {
		check_ajax_referer( 'ai_search_vector_migration' );
		if ( ! current_user_can( 'manage_options' ) ) {
			wp_send_json_error( [ 'message' => __( 'Not allowed.', 'ai-search' ) ], 403 );
		}
		if ( ! get_option( self::OPTION_STARTED_AT ) ) {
			update_option( self::OPTION_STARTED_AT, time(), false );
		}
		delete_transient( self::TRANSIENT_TOTAL );
		wp_send_json_success( self::progress() );
	}

	public static function ajax_batch(): void {
		check_ajax_referer( 'ai_search_vector_migration' );
		if ( ! current_user_can( 'manage_options' ) ) {
			wp_send_json_error( [ 'message' => __( 'Not allowed.', 'ai-search' ) ], 403 );
		}
		$result   = self::migrate_batch();
		$progress = self::progress();
		wp_send_json_success( array_merge( $result, [ 'progress' => $progress ] ) );
	}

	public static function ajax_finalize(): void {
		check_ajax_referer( 'ai_search_vector_migration' );
		if ( ! current_user_can( 'manage_options' ) ) {
			wp_send_json_error( [ 'message' => __( 'Not allowed.', 'ai-search' ) ], 403 );
		}
		$ok = self::finalize();
		if ( ! $ok ) {
			wp_send_json_error( [
				'message'  => __( 'Still have pending rows. Run more batches first.', 'ai-search' ),
				'progress' => self::progress(),
			], 409 );
		}
		wp_send_json_success( self::progress() );
	}

	public static function ajax_dismiss(): void {
		check_ajax_referer( 'ai_search_vector_migration' );
		if ( ! current_user_can( 'manage_options' ) ) {
			wp_send_json_error( [ 'message' => __( 'Not allowed.', 'ai-search' ) ], 403 );
		}
		update_option( self::OPTION_DISMISSED, time() + ( 30 * DAY_IN_SECONDS ), false );
		wp_send_json_success();
	}

	public static function is_dismissed(): bool {
		return (int) get_option( self::OPTION_DISMISSED, 0 ) > time();
	}

	/* -------------------------------------------------------------
	 * Hybrid auto-migration — daily cron that catches every install
	 * silently (no admin interaction needed), with the UI card still
	 * available for admins who want to force completion right now.
	 * ----------------------------------------------------------- */

	/**
	 * Schedule the daily cron if not already scheduled. Called from the
	 * version-upgrade path in the main plugin bootstrap.
	 */
	public static function schedule_cron(): void {
		if ( ! wp_next_scheduled( self::CRON_HOOK ) ) {
			// Fire the first run almost immediately (60s) so the merchant
			// sees progress on the day of the plugin update, not tomorrow.
			wp_schedule_event( time() + MINUTE_IN_SECONDS, 'daily', self::CRON_HOOK );
		}
	}

	/**
	 * Remove the recurring job — called once we're on v2 or when the
	 * admin manually completes via the UI card.
	 */
	public static function unschedule_cron(): void {
		$next = wp_next_scheduled( self::CRON_HOOK );
		while ( $next ) {
			wp_unschedule_event( $next, self::CRON_HOOK );
			$next = wp_next_scheduled( self::CRON_HOOK );
		}
	}

	/**
	 * One daily cron tick. Runs migrate_batch() repeatedly within a bounded
	 * time budget, finalising when zero posts remain. Lock prevents a slow
	 * cron tick from being invoked twice in parallel (overlapping spawn
	 * requests on busy sites).
	 */
	public static function cron_tick(): void {
		// Guard: already migrated.
		if ( self::is_v2_active() ) {
			self::unschedule_cron();
			return;
		}

		// Soft lock to avoid double-firing. 60s is longer than our budget.
		if ( false !== get_transient( self::CRON_LOCK_TRANSIENT ) ) {
			return;
		}
		set_transient( self::CRON_LOCK_TRANSIENT, 1, 60 );

		// Give PHP room to breathe — this is a cron spawn, not a user
		// request. Cap at 20s just in case the host has a tight hard limit.
		if ( function_exists( 'set_time_limit' ) ) {
			@set_time_limit( 20 );
		}

		$deadline = microtime( true ) + self::CRON_BUDGET_SECONDS;
		while ( microtime( true ) < $deadline ) {
			$r = self::migrate_batch();
			if ( 0 === $r['remaining'] ) {
				if ( self::finalize() ) {
					self::unschedule_cron();
					set_transient( self::DONE_NOTICE_KEY, time(), 7 * DAY_IN_SECONDS );
				}
				break;
			}
			if ( 0 === $r['processed'] ) {
				// Nothing moved — avoid tight loop on a stuck batch.
				break;
			}
		}

		delete_transient( self::CRON_LOCK_TRANSIENT );
	}

	/**
	 * Called from the main plugin's version-upgrade path. Idempotent:
	 * safe to call on every load. Schedules the cron if (a) we're still
	 * on v1 AND (b) there's actually work to do.
	 */
	public static function ensure_scheduled(): void {
		if ( self::is_v2_active() ) {
			self::unschedule_cron();
			return;
		}
		if ( self::count_pending() > 0 ) {
			self::schedule_cron();
		}
	}

	/**
	 * One-off success notice shown to the admin after the cron finished
	 * the migration in the background. Dismissed on first view.
	 */
	public static function maybe_render_done_notice(): void {
		if ( ! current_user_can( 'manage_options' ) ) {
			return;
		}
		$fired_at = (int) get_transient( self::DONE_NOTICE_KEY );
		if ( $fired_at <= 0 ) {
			return;
		}
		delete_transient( self::DONE_NOTICE_KEY );
		printf(
			'<div class="notice notice-success is-dismissible"><p><strong>%s</strong> %s</p></div>',
			esc_html__( 'AI Search index upgraded.', 'ai-search' ),
			esc_html__( 'Your search vectors were converted to the new packed format in the background. Searches now run on a smaller, faster index — no action needed from you.', 'ai-search' )
		);
	}
}