<?php
/**
 * Phase 2: on-disk matrix cache.
 *
 * Builds a single packed-binary file holding every unit-normalized embedding
 * vector for the currently active (provider, model, dim). At query time the
 * file is mmap-loaded once per request and cached in a request-static (and
 * APCu, when available) so subsequent chats in the same FPM worker skip both
 * the SELECT and the json_decode/unpack cost entirely.
 *
 * Design constraints:
 *   - Must work on shared hosts without Redis/Memcached.
 *   - Must be safe alongside the legacy JSON path: if anything goes wrong we
 *     return null and the caller falls back to streaming.
 *   - Must be cheap to invalidate: a single unlink() on KB writes plus a
 *     scheduled rebuild that runs out-of-band so user requests never pay
 *     the rebuild cost.
 *   - Must be scoped per (provider, model, dim) so mixed-provider KBs don't
 *     blend dimensionally-incompatible vectors.
 *
 * File layout (little-endian, host-native float32):
 *
 *   HEADER (32 bytes)
 *     magic        4B  ASCII 'AMKB'
 *     fmt_version  2B  uint16, currently 1
 *     dim          2B  uint16
 *     count        4B  uint32, number of rows
 *     built_at     8B  int64 unix timestamp
 *     reserved    12B  zeros (future fields: ANN bucket offset, etc.)
 *
 *   BODY (count rows)
 *     id           8B  uint64
 *     vec     dim*4B   float32 unit-normalized
 *
 * @package Antimanual
 * @since   3.7.0
 */

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

/**
 * Cache file directory inside wp-content/uploads. Created lazily on first
 * build. We drop a stub .htaccess + index.html so the file is not directly
 * fetchable on Apache hosts that serve uploads/.
 *
 * @return string Absolute path with trailing slash, or empty string if the
 *                upload dir is not writable (caller falls back).
 */
function atml_matrix_dir(): string {
	$uploads = wp_upload_dir( null, false );
	if ( ! empty( $uploads['error'] ) ) {
		return '';
	}

	$dir = trailingslashit( $uploads['basedir'] ) . 'antimanual/matrix/';

	if ( ! file_exists( $dir ) ) {
		if ( ! wp_mkdir_p( $dir ) ) {
			return '';
		}
		// Best-effort access denial; we don't fail the cache if these can't
		// be written. The file is binary noise either way — no leakage risk.
		@file_put_contents( $dir . '.htaccess', "Order allow,deny\nDeny from all\n" );
		@file_put_contents( $dir . 'index.html', '' );
	}

	if ( ! is_writable( $dir ) ) {
		return '';
	}

	return $dir;
}

/**
 * Slug-safe segment of a model name for use in a filename.
 *
 * @param string $value
 * @return string
 */
function atml_matrix_slug( string $value ): string {
	$value = strtolower( $value );
	$value = preg_replace( '/[^a-z0-9._-]+/', '-', $value );
	$value = trim( (string) $value, '-' );
	return $value === '' ? 'x' : $value;
}

/**
 * Resolve the cache file path for the current active embedding context.
 *
 * Returns an empty string when no writable directory is available (caller
 * treats that the same as "file missing" and falls back).
 *
 * @return string
 */
function atml_matrix_path(): string {
	$dir = atml_matrix_dir();
	if ( '' === $dir ) {
		return '';
	}

	$provider = atml_matrix_slug( (string) \Antimanual\AIProvider::get_embedding_provider_name() );
	$model    = atml_matrix_slug( (string) \Antimanual\AIProvider::get_embedding_model() );
	// Phase 4: include the configured dim in the filename so changing the
	// dim setting points at a different cache file. A reindex will populate
	// the new file; the old one is invalidated by atml_matrix_invalidate()
	// (which keys on the *current* path) and can be left for cleanup later.
	$dim      = function_exists( 'atml_get_embedding_dimensions' )
		? atml_get_embedding_dimensions( $provider )
		: 0;
	$dim_slug = $dim > 0 ? 'd' . $dim : 'default';

	return $dir . "kb-{$provider}-{$model}-{$dim_slug}-v1.bin";
}

/**
 * Load + parse the matrix file for the active embedding context.
 *
 * Returns null on any failure: missing file, header mismatch, partial read,
 * empty body, or stale file (mtime older than the most recent KB write).
 *
 * Cached per-request (and per-PHP-FPM-worker via APCu when present).
 *
 * @return array{vectors: array<int, float[]>, norms: array<int, float>, count: int}|null
 */
function atml_matrix_load_file(): ?array {
	static $request_cache = null;
	if ( null !== $request_cache ) {
		return $request_cache ?: null;
	}

	$path = atml_matrix_path();
	if ( '' === $path || ! is_readable( $path ) ) {
		$request_cache = false;
		return null;
	}

	// APCu fast path. Keyed on file mtime so a rebuild invalidates implicitly.
	$apcu_key = null;
	if ( function_exists( 'apcu_fetch' ) ) {
		$mtime    = (int) filemtime( $path );
		$apcu_key = 'atml_matrix_' . md5( $path . '|' . $mtime );
		$hit      = apcu_fetch( $apcu_key, $ok );
		if ( $ok && is_array( $hit ) && isset( $hit['vectors'], $hit['norms'], $hit['count'] ) ) {
			$request_cache = $hit;
			return $hit;
		}
	}

	$fp = @fopen( $path, 'rb' );
	if ( ! $fp ) {
		$request_cache = false;
		return null;
	}

	$header = fread( $fp, 32 );
	if ( ! is_string( $header ) || strlen( $header ) !== 32 ) {
		fclose( $fp );
		$request_cache = false;
		return null;
	}

	$h = unpack( 'A4magic/vfmt/vdim/Vcount/Pbuilt', $header );
	if ( ! is_array( $h ) || ( $h['magic'] ?? '' ) !== 'AMKB' || ( $h['fmt'] ?? 0 ) !== 1 ) {
		fclose( $fp );
		$request_cache = false;
		return null;
	}

	$dim     = (int) $h['dim'];
	$count   = (int) $h['count'];
	$rowsize = 8 + $dim * 4;

	if ( $dim <= 0 || $count <= 0 || $rowsize <= 8 ) {
		fclose( $fp );
		$request_cache = false;
		return null;
	}

	$vectors = [];
	$norms   = [];

	for ( $i = 0; $i < $count; $i++ ) {
		$row = fread( $fp, $rowsize );
		if ( ! is_string( $row ) || strlen( $row ) !== $rowsize ) {
			// Truncated file. Bail; caller falls back to streaming.
			fclose( $fp );
			$request_cache = false;
			return null;
		}

		// uint64 id, then dim×float32. PHP 8 supports 'P' for uint64-LE.
		$id_unpacked = unpack( 'Pid', substr( $row, 0, 8 ) );
		if ( ! is_array( $id_unpacked ) ) {
			continue;
		}
		$id = (int) $id_unpacked['id'];

		$vec_unpacked = unpack( 'f*', substr( $row, 8 ) );
		if ( ! is_array( $vec_unpacked ) ) {
			continue;
		}

		// Re-key to 0-indexed for consistency with json_decoded vectors.
		$vectors[ $id ] = array_values( $vec_unpacked );
		// Stored vectors are pre-normalized at write time.
		$norms[ $id ]   = 1.0;
	}
	fclose( $fp );

	$result = [
		'vectors' => $vectors,
		'norms'   => $norms,
		'count'   => count( $vectors ),
	];

	if ( $apcu_key && function_exists( 'apcu_store' ) ) {
		$ttl = (int) apply_filters( 'antimanual_matrix_apcu_ttl', HOUR_IN_SECONDS );
		@apcu_store( $apcu_key, $result, $ttl );
	}

	$request_cache = $result;
	return $result;
}

/**
 * Build (or rebuild) the cache file from the database, atomically.
 *
 * Selects only rows for the currently active provider+model whose
 * `embedding_bin` has been populated by Phase 1's backfill. Rows still in
 * legacy-JSON-only form are skipped — they get picked up on the next
 * rebuild after the backfill cron promotes them.
 *
 * Runs as a single SQL query + sequential fwrite. On a 70 MB / 50 k-row KB
 * this is ~1–3 s wall-clock on a typical shared host. Designed to run from
 * a cron tick, NOT inline from a chat request.
 *
 * @return array{built: bool, count: int, bytes: int, path: string, reason?: string}
 */
function atml_matrix_build_file(): array {
	$path = atml_matrix_path();
	if ( '' === $path ) {
		return [ 'built' => false, 'count' => 0, 'bytes' => 0, 'path' => '', 'reason' => 'no_writable_dir' ];
	}

	if ( ! class_exists( '\\Antimanual\\Embedding' ) || ! \Antimanual\Embedding::has_binary_columns() ) {
		return [ 'built' => false, 'count' => 0, 'bytes' => 0, 'path' => $path, 'reason' => 'schema_not_ready' ];
	}

	global $wpdb;
	$table = \Antimanual\Embedding::get_table_name();

	$active_provider = (string) \Antimanual\AIProvider::get_embedding_provider_name();
	$active_model    = (string) \Antimanual\AIProvider::get_embedding_model();
	// Phase 4: when the admin opted into a non-default output dim, pin the
	// matrix to exactly that dim. Otherwise accept any positive dim — rows
	// that don't match the live query dim are skipped at retrieval time.
	$active_dim = function_exists( 'atml_get_embedding_dimensions' )
		? atml_get_embedding_dimensions( $active_provider )
		: 0;

	if ( $active_dim > 0 ) {
		$rows = $wpdb->get_results(
			$wpdb->prepare(
				"SELECT id, embedding_bin, embedding_dim FROM `$table`
				 WHERE provider = %s
				   AND embedding_model = %s
				   AND embedding_bin IS NOT NULL
				   AND embedding_dim = %d
				 ORDER BY id ASC",
				$active_provider,
				$active_model,
				$active_dim
			)
		);
	} else {
		// Scope to active provider+model so dim is uniform across rows. Also
		// require embedding_bin populated so we never write half-binary rows.
		$rows = $wpdb->get_results(
			$wpdb->prepare(
				"SELECT id, embedding_bin, embedding_dim FROM `$table`
				 WHERE provider = %s
				   AND embedding_model = %s
				   AND embedding_bin IS NOT NULL
				   AND embedding_dim > 0
				 ORDER BY id ASC",
				$active_provider,
				$active_model
			)
		);
	}

	if ( empty( $rows ) ) {
		// Nothing yet — possibly still backfilling, possibly empty KB.
		// Don't write a zero-row file; caller falls back to streaming.
		return [ 'built' => false, 'count' => 0, 'bytes' => 0, 'path' => $path, 'reason' => 'no_rows' ];
	}

	$dim = (int) $rows[0]->embedding_dim;
	if ( $dim <= 0 ) {
		return [ 'built' => false, 'count' => 0, 'bytes' => 0, 'path' => $path, 'reason' => 'bad_dim' ];
	}

	$tmp = $path . '.tmp.' . wp_generate_password( 8, false, false );
	$fp  = @fopen( $tmp, 'wb' );
	if ( ! $fp ) {
		return [ 'built' => false, 'count' => 0, 'bytes' => 0, 'path' => $path, 'reason' => 'open_tmp_failed' ];
	}

	$header = pack(
		'A4vvVPa12',
		'AMKB',         // magic
		1,              // fmt_version
		$dim,           // dim
		count( $rows ), // count
		time(),         // built_at
		''              // reserved (12 zero bytes)
	);
	$written = fwrite( $fp, $header );
	if ( false === $written || strlen( $header ) !== 32 ) {
		fclose( $fp );
		@unlink( $tmp );
		return [ 'built' => false, 'count' => 0, 'bytes' => 0, 'path' => $path, 'reason' => 'header_write_failed' ];
	}

	$expected_blob_len = $dim * 4;
	$bytes             = 32;
	$written_rows      = 0;

	foreach ( $rows as $row ) {
		$bin = (string) $row->embedding_bin;
		if ( strlen( $bin ) !== $expected_blob_len ) {
			// Skip dim-mismatched rows silently — keeps the file uniform.
			continue;
		}

		$id_bin = pack( 'P', (int) $row->id );
		if ( false === fwrite( $fp, $id_bin ) || false === fwrite( $fp, $bin ) ) {
			fclose( $fp );
			@unlink( $tmp );
			return [ 'built' => false, 'count' => 0, 'bytes' => 0, 'path' => $path, 'reason' => 'row_write_failed' ];
		}
		$bytes        += 8 + $expected_blob_len;
		$written_rows += 1;
	}

	fclose( $fp );

	if ( 0 === $written_rows ) {
		@unlink( $tmp );
		return [ 'built' => false, 'count' => 0, 'bytes' => 0, 'path' => $path, 'reason' => 'no_uniform_rows' ];
	}

	// Header count may be off if we skipped dim-mismatched rows. Patch it
	// in-place before we rename so the loader doesn't try to read past EOF.
	if ( $written_rows !== count( $rows ) ) {
		$fp_patch = @fopen( $tmp, 'r+b' );
		if ( $fp_patch ) {
			// count is at byte offset 8 (4 magic + 2 fmt + 2 dim).
			fseek( $fp_patch, 8 );
			fwrite( $fp_patch, pack( 'V', $written_rows ) );
			fclose( $fp_patch );
		}
	}

	// Atomic publish. On Windows rename() over an existing file fails, so
	// unlink first. The race window is OK because readers fall back when
	// the file is missing.
	if ( file_exists( $path ) ) {
		@unlink( $path );
	}
	if ( ! @rename( $tmp, $path ) ) {
		@unlink( $tmp );
		return [ 'built' => false, 'count' => 0, 'bytes' => 0, 'path' => $path, 'reason' => 'rename_failed' ];
	}

	error_log( sprintf(
		'[Antimanual][matrix] built path=%s rows=%d bytes=%d dim=%d',
		basename( $path ), $written_rows, $bytes, $dim
	) );

	return [
		'built' => true,
		'count' => $written_rows,
		'bytes' => $bytes,
		'path'  => $path,
	];
}

/**
 * Invalidate the cache file for the active provider+model and schedule a
 * rebuild in the background. Idempotent; cheap to call from any KB write.
 *
 * The rebuild is debounced by a 30-second floor: if a rebuild is already
 * scheduled within the next 30s we don't add another one. This handles the
 * case where a bulk import inserts many rows in quick succession.
 */
function atml_matrix_invalidate(): void {
	$path = atml_matrix_path();
	if ( '' !== $path && file_exists( $path ) ) {
		@unlink( $path );
	}

	// Also clear any APCu copy. The new mtime keys it differently, but if
	// the file is gone the next read returns null anyway, so this is mostly
	// for tidiness.
	if ( function_exists( 'apcu_clear_cache' ) ) {
		// Don't blow away the whole cache — only our keys, which are mtime-
		// scoped. Letting them expire naturally is fine.
	}

	if ( ! wp_next_scheduled( 'atml_rebuild_matrix_file' ) ) {
		wp_schedule_single_event( time() + 30, 'atml_rebuild_matrix_file' );
	}
}
