<?php
/**
 * Builder-side RAG indexer.
 *
 * Embeds builder-owned content (header/footer templates, components,
 * registered shortcodes, doc/ARCHITECTURE.md) so the AI chat can retrieve
 * concrete examples instead of relying on a static system prompt dump.
 *
 * @package Antimanual_Builder
 * @since   1.6.0
 */

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

class AMBLD_Rag_Indexer {

	const MAX_CHUNK_CHARS = 1800;

	/** wp_options keys driving the seed bootstrap cron. */
	const OPT_SEED_PENDING       = 'ambld_seed_pending';
	const OPT_SEED_COMPLETED_AT  = 'ambld_seed_completed_at';
	const OPT_SEED_RETRY_COUNT   = 'ambld_seed_retry_count';
	const OPT_SEED_LAST_ERROR    = 'ambld_seed_last_error';
	const CRON_HOOK              = 'ambld_seed_embed_run';
	const MAX_SEED_RETRIES       = 24;

	/**
	 * Provider → expected embedding model + dimension for pre-shipped vector
	 * files. Used by both the build script (so it writes the canonical
	 * filename) and the runtime loader (so it knows which file to look for
	 * given the user's active embedding provider).
	 *
	 * Filename convention: `{provider}-{model}-{dim}.{json,bin}` under
	 * `data/seed/embeddings/`. Slug derivation lives in
	 * `shipped_vectors_slug()` and is canonical — DO NOT diverge between
	 * the build script and the loader.
	 */
	const SHIPPED_PROVIDERS = array(
		'openai' => array( 'model' => 'text-embedding-3-small', 'dim' => 1536 ),
		// Gemini's gemini-embedding-001 returns 3072-dim vectors by default.
		// Parent plugin (Gemini::generate_embedding) doesn't pass
		// outputDimensionality, so query-time vectors are 3072 too — shipped
		// vectors MUST match this dim or the retriever's dim-mismatch skip
		// will silently drop every Gemini seed row.
		'gemini' => array( 'model' => 'gemini-embedding-001',   'dim' => 3072 ),
	);

	/** Subdirectory under `data/seed/` that holds .json pre-baked vector files (base64 float32 inline). */
	const SHIPPED_VECTORS_REL = 'data/seed/embeddings/';

	public function register() {
		add_action( 'save_post_ambld_template', array( $this, 'index_template' ), 20, 1 );
		add_action( 'save_post_ambld_component', array( $this, 'index_component' ), 20, 1 );
		add_action( 'before_delete_post', array( $this, 'delete_for_post' ), 10, 1 );
	}

	/**
	 * Whether the parent Antimanual plugin's embedding service is reachable.
	 * Without it we cannot vectorize content; the indexer is a no-op until
	 * the dependency is active.
	 */
	public static function can_embed() {
		return class_exists( '\Antimanual\AIProvider' )
			&& method_exists( '\Antimanual\AIProvider', 'generate_embedding' );
	}

	private static function embed( $text ) {
		if ( ! self::can_embed() ) {
			return null;
		}
		$vector = \Antimanual\AIProvider::generate_embedding( (string) $text );
		if ( is_wp_error( $vector ) || empty( $vector ) || ! is_array( $vector ) ) {
			return null;
		}
		return array_map( 'floatval', $vector );
	}

	private static function embedding_model_label() {
		if ( class_exists( '\Antimanual\AIProvider' ) && method_exists( '\Antimanual\AIProvider', 'get_embedding_provider_name' ) ) {
			return (string) \Antimanual\AIProvider::get_embedding_provider_name();
		}
		return 'unknown';
	}

	/**
	 * Split a string into chunks bounded by MAX_CHUNK_CHARS, preferring
	 * paragraph and sentence boundaries so we don't slice a CSS rule or
	 * sentence in half mid-token.
	 */
	public static function chunk_text( $text, $max_chars = self::MAX_CHUNK_CHARS ) {
		$text = trim( (string) $text );
		if ( '' === $text ) {
			return array();
		}
		if ( strlen( $text ) <= $max_chars ) {
			return array( $text );
		}

		$paragraphs = preg_split( '/\n\s*\n/', $text );
		$chunks     = array();
		$buffer     = '';

		foreach ( $paragraphs as $para ) {
			if ( strlen( $buffer ) + strlen( $para ) + 2 > $max_chars && '' !== $buffer ) {
				$chunks[] = trim( $buffer );
				$buffer   = '';
			}
			if ( strlen( $para ) > $max_chars ) {
				// Hard wrap oversize paragraphs at word boundaries.
				$wrapped = wordwrap( $para, $max_chars, "\n---SPLIT---\n", true );
				foreach ( explode( "\n---SPLIT---\n", $wrapped ) as $sub ) {
					$chunks[] = trim( $sub );
				}
				continue;
			}
			$buffer .= ( '' === $buffer ? '' : "\n\n" ) . $para;
		}
		if ( '' !== trim( $buffer ) ) {
			$chunks[] = trim( $buffer );
		}
		return array_values( array_filter( $chunks, 'strlen' ) );
	}

	public function index_template( $post_id ) {
		if ( wp_is_post_revision( $post_id ) || wp_is_post_autosave( $post_id ) ) {
			return;
		}
		$post = get_post( $post_id );
		if ( ! $post || 'ambld_template' !== $post->post_type ) {
			return;
		}
		$kind  = (string) get_post_meta( $post_id, '_ambld_template_kind', true );
		$title = $post->post_title ? $post->post_title : sprintf( 'Template #%d', $post_id );
		$body  = (string) $post->post_content;
		$head  = sprintf( "Template: %s\nKind: %s\n\n", $title, $kind ?: 'generic' );
		$this->index_text( 'template', (string) $post_id, $title, $head . $body );
	}

	public function index_component( $post_id ) {
		if ( wp_is_post_revision( $post_id ) || wp_is_post_autosave( $post_id ) ) {
			return;
		}
		$post = get_post( $post_id );
		if ( ! $post || 'ambld_component' !== $post->post_type ) {
			return;
		}
		$title = $post->post_title ? $post->post_title : sprintf( 'Component #%d', $post_id );
		$body  = (string) $post->post_content;
		$head  = sprintf( "Component: %s\n\n", $title );
		$this->index_text( 'component', (string) $post_id, $title, $head . $body );
	}

	public function index_shortcodes() {
		global $shortcode_tags;
		if ( ! is_array( $shortcode_tags ) ) {
			return;
		}
		foreach ( array_keys( $shortcode_tags ) as $tag ) {
			$tag      = (string) $tag;
			$content  = sprintf( "Shortcode: [%s]\nUsage: insert via the shortcode block. Output is rendered server-side via do_shortcode.", $tag );
			$this->index_text( 'shortcode', $tag, '[' . $tag . ']', $content );
		}
	}

	public function index_architecture_doc() {
		$path = AMBLD_PLUGIN_DIR . 'doc/ARCHITECTURE.md';
		if ( ! file_exists( $path ) ) {
			return;
		}
		$body = file_get_contents( $path ); // phpcs:ignore WordPress.WP.AlternativeFunctions
		if ( ! is_string( $body ) || '' === trim( $body ) ) {
			return;
		}
		$this->index_text( 'doc', 'architecture', 'ARCHITECTURE.md', $body );
	}

	public function delete_for_post( $post_id ) {
		$post = get_post( $post_id );
		if ( ! $post ) {
			return;
		}
		if ( 'ambld_template' === $post->post_type ) {
			AMBLD_Rag_Store::delete_source( 'template', (string) $post_id );
		} elseif ( 'ambld_component' === $post->post_type ) {
			AMBLD_Rag_Store::delete_source( 'component', (string) $post_id );
		}
	}

	/**
	 * Embed and store all chunks of a piece of content.
	 *
	 * Skips silently when the parent embedding service isn't available, so
	 * activations on a fresh site don't fatal.
	 */
	private function index_text( $source_type, $source_id, $label, $text ) {
		if ( ! self::can_embed() ) {
			return;
		}
		$chunks   = self::chunk_text( $text );
		$prepared = array();
		foreach ( $chunks as $chunk ) {
			$vec = self::embed( $chunk );
			if ( null === $vec ) {
				continue;
			}
			$prepared[] = array(
				'content'   => $chunk,
				'embedding' => $vec,
				'label'     => $label,
			);
		}
		if ( empty( $prepared ) ) {
			return;
		}
		AMBLD_Rag_Store::replace_source( $source_type, $source_id, $prepared, self::embedding_model_label() );
	}

	/**
	 * Reindex all builder-side sources. Invoked by the WP-CLI command
	 * `wp amb rag reindex` and by activation when feasible.
	 */
	public function reindex_all() {
		if ( ! self::can_embed() ) {
			return new WP_Error( 'no_embedder', 'Antimanual embedding service unavailable.' );
		}
		$counts = array( 'templates' => 0, 'components' => 0, 'shortcodes' => 0, 'doc' => 0, 'seed' => 0 );

		$templates = get_posts( array( 'post_type' => 'ambld_template', 'posts_per_page' => -1, 'post_status' => 'any' ) );
		foreach ( $templates as $t ) {
			$this->index_template( $t->ID );
			$counts['templates']++;
		}

		$components = get_posts( array( 'post_type' => 'ambld_component', 'posts_per_page' => -1, 'post_status' => 'any' ) );
		foreach ( $components as $c ) {
			$this->index_component( $c->ID );
			$counts['components']++;
		}

		$this->index_shortcodes();
		$counts['shortcodes'] = is_array( $GLOBALS['shortcode_tags'] ?? null ) ? count( $GLOBALS['shortcode_tags'] ) : 0;

		$this->index_architecture_doc();
		$counts['doc'] = 1;

		$seed = $this->seed_corpus();
		if ( ! is_wp_error( $seed ) && isset( $seed['embedded'] ) ) {
			$counts['seed'] = (int) $seed['embedded'];
		}

		return $counts;
	}

	/**
	 * Index the shipped seed corpus under data/seed/.
	 *
	 * Strategy (in order):
	 *   1. **Pre-baked vector path (no API cost).** If the maintainer has
	 *      shipped `.json` + `.bin` files for the user's active embedding
	 *      provider, load each seed file's vectors directly from disk and
	 *      write them to the store. This is the default path for the
	 *      ~95% of installs running OpenAI or Gemini and incurs zero API
	 *      cost.
	 *   2. **API-embed fallback.** For any seed file NOT covered by the
	 *      shipped vectors (corpus edited since the maintainer's build,
	 *      exotic provider with no shipped file, drift between body hash
	 *      and vector), embed via the parent Antimanual provider as
	 *      before. Each source's previous rows are replaced wholesale.
	 *   3. **No-op.** Files already indexed at the current content_hash
	 *      are skipped via the existing `content_hash_exists()` probe.
	 *
	 * Returns a stats tuple useful for admin diagnostics and WP-CLI.
	 *
	 * @return array{embedded:int, skipped:int, errors:int, installed_from_vectors:int}|WP_Error
	 */
	public function seed_corpus() {
		if ( ! class_exists( 'AMBLD_Rag_Seed_Loader' ) ) {
			return new WP_Error( 'no_loader', 'Seed loader not available.' );
		}

		$records = AMBLD_Rag_Seed_Loader::load_all();
		$result  = array(
			'embedded'              => 0,
			'skipped'               => 0,
			'errors'                => 0,
			'installed_from_vectors' => 0,
		);
		if ( empty( $records ) ) {
			return $result;
		}

		error_log( sprintf(
			'[AMB-RAG] seed_corpus: START | records_in_corpus=%d',
			count( $records )
		) );

		// --- Step 1: install pre-baked shipped vectors when available. ---
		// We do this even when the embedder isn't configured — reading vectors
		// from disk requires no API key. This is the entire point of shipping
		// pre-computed embeddings: a fresh install with the right provider
		// gets a working RAG corpus without ever calling out to OpenAI / Gemini.
		$shipped = $this->load_shipped_vectors_for_active_provider();
		if ( null === $shipped ) {
			error_log( '[AMB-RAG] seed_corpus: shipped vectors NOT loaded (provider unsupported, files missing, or manifest invalid) — falling through to API-embed path' );
		} else {
			error_log( sprintf(
				'[AMB-RAG] seed_corpus: shipped vectors LOADED for model=%s dim=%d sources=%d',
				$shipped['model'],
				$shipped['dim'],
				count( $shipped['by_source'] )
			) );
		}
		if ( is_array( $shipped ) ) {
			$remaining = array();
			foreach ( $records as $record ) {
				$entry = isset( $shipped['by_source'][ $record['source_id'] ] )
					? $shipped['by_source'][ $record['source_id'] ]
					: null;

				// Match only when the body hash agrees. If the seed file
				// has been edited since the maintainer ran the build script,
				// the shipped chunks describe stale content — fall through
				// to API embed for this source.
				if ( $entry && $entry['body_hash'] === $record['content_hash'] && ! empty( $entry['chunks'] ) ) {
					if ( AMBLD_Rag_Store::content_hash_exists( $record['source_type'], $record['source_id'], $record['content_hash'] ) ) {
						$result['skipped']++;
						continue;
					}
					AMBLD_Rag_Store::replace_source(
						$record['source_type'],
						$record['source_id'],
						$entry['chunks'],
						(string) $shipped['model']
					);
					$result['installed_from_vectors']++;
					continue;
				}
				$remaining[] = $record;
			}
			$records = $remaining;
		}

		// --- Step 2: API-embed fallback for uncovered records. ---
		if ( empty( $records ) ) {
			return $result;
		}
		if ( ! self::can_embed() ) {
			// Shipped vectors took whatever they could; the rest waits for
			// the embedder to come online. Not an error — leave pending so
			// the cron retries.
			return $result;
		}

		foreach ( $records as $record ) {
			if ( AMBLD_Rag_Store::content_hash_exists( $record['source_type'], $record['source_id'], $record['content_hash'] ) ) {
				$result['skipped']++;
				continue;
			}

			$chunks   = self::chunk_text( $record['content'] );
			$prepared = array();
			foreach ( $chunks as $chunk ) {
				$vec = self::embed( $chunk );
				if ( null === $vec ) {
					$result['errors']++;
					continue 2; // skip this whole source on any chunk failure — we'll retry next pass
				}
				$prepared[] = array(
					'content'   => $chunk,
					'embedding' => $vec,
					'label'     => $record['source_label'],
				);
			}
			if ( empty( $prepared ) ) {
				continue;
			}
			AMBLD_Rag_Store::replace_source(
				$record['source_type'],
				$record['source_id'],
				$prepared,
				self::embedding_model_label()
			);
			$result['embedded']++;
		}

		error_log( sprintf(
			'[AMB-RAG] seed_corpus: DONE | installed_from_vectors=%d embedded_via_api=%d skipped=%d errors=%d',
			$result['installed_from_vectors'],
			$result['embedded'],
			$result['skipped'],
			$result['errors']
		) );

		return $result;
	}

	/**
	 * Canonical filename slug for a shipped vector file. Used by both the
	 * build script and the loader so they agree byte-for-byte on the path.
	 *
	 * @param  string $provider Provider slug (lowercase, e.g. "openai").
	 * @param  string $model    Model name (lowercase, hyphenated).
	 * @param  int    $dim      Embedding dimension.
	 * @return string Slug without extension (e.g. "openai-text-embedding-3-small-1536").
	 */
	public static function shipped_vectors_slug( $provider, $model, $dim ) {
		$provider = strtolower( trim( (string) $provider ) );
		$model    = strtolower( trim( (string) $model ) );
		$dim      = (int) $dim;
		if ( '' === $provider || '' === $model || $dim <= 0 ) {
			return '';
		}
		$clean_model = preg_replace( '/[^a-z0-9\-]+/', '-', $model );
		$clean_model = trim( (string) $clean_model, '-' );
		return sprintf( '%s-%s-%d', $provider, $clean_model, $dim );
	}

	/**
	 * Absolute path of the shipped vectors directory.
	 */
	public static function shipped_vectors_dir() {
		return AMBLD_PLUGIN_DIR . self::SHIPPED_VECTORS_REL;
	}

	/**
	 * Load and decode shipped vectors for the active embedding provider.
	 *
	 * Returns null when:
	 *   - The active provider has no `SHIPPED_PROVIDERS` mapping.
	 *   - The expected `.json` / `.bin` files don't exist on disk.
	 *   - The manifest is invalid or the dimension doesn't match.
	 *
	 * Returning null is NOT an error — it just means we fall back to the
	 * API-embed path.
	 *
	 * @return array{model:string, dim:int, by_source: array<string, array{source_type:string, body_hash:string, chunks: array<int, array{content:string, embedding:array<int,float>, label:string}>}>}|null
	 */
	public function load_shipped_vectors_for_active_provider() {
		$provider_slug = self::embedding_model_label();
		$meta          = isset( self::SHIPPED_PROVIDERS[ $provider_slug ] )
			? self::SHIPPED_PROVIDERS[ $provider_slug ]
			: null;
		if ( ! $meta ) {
			return null;
		}
		return $this->load_shipped_vectors( $provider_slug, $meta['model'], (int) $meta['dim'] );
	}

	/**
	 * Load shipped vectors for a specific (provider, model, dim) tuple.
	 *
	 * Reads `<slug>.json` where each chunk carries an inline `vec` field
	 * (base64-encoded little-endian float32 blob, `dim * 4` bytes when
	 * decoded). Inlining the vectors keeps the corpus to a single text file
	 * — plugin-check rejects bare `.bin` payloads as application files.
	 *
	 * @param  string $provider Provider slug.
	 * @param  string $model    Model name.
	 * @param  int    $dim      Expected embedding dimension.
	 * @return array|null Indexed structure or null on any failure.
	 */
	public function load_shipped_vectors( $provider, $model, $dim ) {
		$slug = self::shipped_vectors_slug( $provider, $model, $dim );
		if ( '' === $slug ) {
			return null;
		}
		$dir       = self::shipped_vectors_dir();
		$json_path = $dir . $slug . '.json';

		if ( ! is_readable( $json_path ) ) {
			return null;
		}

		$json_raw = @file_get_contents( $json_path ); // phpcs:ignore WordPress.WP.AlternativeFunctions, WordPress.PHP.NoSilencedErrors
		$manifest = is_string( $json_raw ) ? json_decode( $json_raw, true ) : null;
		if ( ! is_array( $manifest ) || empty( $manifest['sources'] ) ) {
			return null;
		}
		if ( (int) ( $manifest['dim'] ?? 0 ) !== (int) $dim ) {
			return null;
		}

		$vec_bytes = $dim * 4;
		$by_source = array();
		foreach ( $manifest['sources'] as $source ) {
			if ( empty( $source['source_id'] ) || empty( $source['source_type'] ) || empty( $source['chunks'] ) ) {
				continue;
			}
			$chunks = array();
			foreach ( $source['chunks'] as $chunk ) {
				if ( empty( $chunk['vec'] ) ) {
					continue;
				}
				$blob = base64_decode( (string) $chunk['vec'], true );
				if ( ! is_string( $blob ) || strlen( $blob ) !== $vec_bytes ) {
					continue;
				}
				$unpacked = unpack( 'f*', $blob );
				if ( ! is_array( $unpacked ) || count( $unpacked ) !== $dim ) {
					continue;
				}
				$chunks[] = array(
					'content'   => isset( $chunk['content'] ) ? (string) $chunk['content'] : '',
					'embedding' => array_values( $unpacked ),
					'label'     => isset( $source['label'] ) ? (string) $source['label'] : '',
				);
			}
			if ( empty( $chunks ) ) {
				continue;
			}
			$by_source[ (string) $source['source_id'] ] = array(
				'source_type' => (string) $source['source_type'],
				'body_hash'   => isset( $source['body_hash'] ) ? (string) $source['body_hash'] : '',
				'chunks'      => $chunks,
			);
		}

		return array(
			'model'     => isset( $manifest['model'] ) ? (string) $manifest['model'] : $model,
			'dim'       => $dim,
			'by_source' => $by_source,
		);
	}

	/**
	 * WP-Cron handler that drives the post-activation seed embed.
	 *
	 * Flow:
	 *   1. Always call `seed_corpus()`. The first thing it does is try
	 *      pre-baked shipped vectors — which requires NO embedder. So even
	 *      if the parent Antimanual provider isn't configured, this pass
	 *      may still cover every seed file at zero API cost.
	 *   2. If `seed_corpus()` covered every seed file (sum of installed +
	 *      embedded + skipped == total), clear the pending flag and we're
	 *      done.
	 *   3. Otherwise the embedder is missing AND there are uncovered files.
	 *      Bump the retry counter and reschedule. Capped at MAX_SEED_RETRIES
	 *      so a permanently misconfigured site stops churning after roughly
	 *      a day of hourly retries.
	 *
	 * @return void
	 */
	public static function run_seed_embed_cron() {
		if ( (int) get_option( self::OPT_SEED_PENDING, 0 ) !== 1 ) {
			return; // Nothing pending.
		}

		$indexer = new self();
		$result  = $indexer->seed_corpus();
		if ( is_wp_error( $result ) ) {
			update_option( self::OPT_SEED_LAST_ERROR, $result->get_error_message(), false );
			wp_schedule_single_event( time() + HOUR_IN_SECONDS, self::CRON_HOOK );
			return;
		}

		// Did we cover everything? Count what's in the store under seed_* types.
		$total_seed_files = class_exists( 'AMBLD_Rag_Seed_Loader' ) ? count( AMBLD_Rag_Seed_Loader::load_all() ) : 0;
		$covered          = (int) $result['installed_from_vectors'] + (int) $result['embedded'] + (int) $result['skipped'];

		if ( $covered >= $total_seed_files ) {
			// Full coverage — clear the pending state.
			delete_option( self::OPT_SEED_PENDING );
			delete_option( self::OPT_SEED_RETRY_COUNT );
			delete_option( self::OPT_SEED_LAST_ERROR );
			update_option( self::OPT_SEED_COMPLETED_AT, current_time( 'mysql', true ), false );
			return;
		}

		// Partial coverage — typically because the embedder is unavailable
		// and shipped vectors didn't cover everything (provider not in
		// SHIPPED_PROVIDERS, or seed body edited since the maintainer's
		// build). Stay pending; retry with backoff.
		$retry = (int) get_option( self::OPT_SEED_RETRY_COUNT, 0 );
		if ( $retry >= self::MAX_SEED_RETRIES ) {
			update_option( self::OPT_SEED_LAST_ERROR, sprintf(
				'Gave up after %d retries with partial coverage (covered %d of %d seed files). Admin can re-run via /preferences/rag-reseed.',
				self::MAX_SEED_RETRIES,
				$covered,
				$total_seed_files
			), false );
			return;
		}
		update_option( self::OPT_SEED_RETRY_COUNT, $retry + 1, false );
		wp_schedule_single_event( time() + HOUR_IN_SECONDS, self::CRON_HOOK );
	}

	/**
	 * Schedule a fresh seed embed pass. Used by the activator and by the
	 * admin "Re-run shipped corpus indexing" button.
	 *
	 * Idempotent — sets the pending flag and clears retry state so the next
	 * cron tick runs immediately rather than honoring the backoff window.
	 *
	 * @param int $delay_seconds Optional delay before the cron fires.
	 * @return void
	 */
	public static function schedule_seed_embed( $delay_seconds = 30 ) {
		update_option( self::OPT_SEED_PENDING, 1, false );
		delete_option( self::OPT_SEED_RETRY_COUNT );
		// Drop any previously-scheduled event so a manual re-run isn't queued
		// behind a stale backoff.
		$existing = wp_next_scheduled( self::CRON_HOOK );
		if ( $existing ) {
			wp_unschedule_event( $existing, self::CRON_HOOK );
		}
		wp_schedule_single_event( time() + max( 5, (int) $delay_seconds ), self::CRON_HOOK );
	}
}
