<?php
/**
 * Top-K cosine retrieval over the builder RAG index.
 *
 * In-process scoring against AMBLD_Rag_Store::load_all() — appropriate for
 * builder-scale corpora (hundreds of templates, not millions). For larger
 * scale, swap this for a vector DB without changing callers.
 *
 * @package Antimanual_Builder
 * @since   1.6.0
 */

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

class AMBLD_Rag_Retriever {

	/**
	 * Per-process diagnostic record from the most recent retrieve() call.
	 *
	 * Filled in by retrieve() at every exit point. Read by the harness via
	 * `last_diagnostic()` so it can surface "WHY did this turn return 0
	 * hits?" to the editor. The dim-mismatch case is the most common
	 * silent-failure mode: stored vectors were embedded at one dim, the
	 * live query embedding is at a different dim, and every row gets
	 * skipped during cosine scoring.
	 *
	 * Shape:
	 *   skip_reason: '' | 'empty_query' | 'no_embedder' | 'embed_error'
	 *                | 'embed_empty' | 'zero_norm' | 'no_rows'
	 *                | ''                                     ← success (any hit count)
	 *   query_chars: int
	 *   query_dim:   int                                       ← 0 when skipped before embedding
	 *   rows_loaded: int                                       ← rows pulled from store
	 *   scored:      int                                       ← rows with matching dim
	 *   skipped_dim: int                                       ← rows skipped (mismatched dim)
	 *   hits:        int                                       ← returned top-K count
	 *   elapsed_ms:  int
	 *   error:       string                                    ← embed error message when relevant
	 *
	 * @var array
	 */
	private static $last_diagnostic = array();

	public static function last_diagnostic() {
		return self::$last_diagnostic;
	}

	private static function record_diagnostic( array $patch ) {
		self::$last_diagnostic = array_merge(
			array(
				'skip_reason' => '',
				'query_chars' => 0,
				'query_dim'   => 0,
				'rows_loaded' => 0,
				'scored'      => 0,
				'skipped_dim' => 0,
				'hits'        => 0,
				'elapsed_ms'  => 0,
				'error'       => '',
			),
			$patch
		);
	}

	/**
	 * Embed the user query and return top-K matching chunks.
	 *
	 * @param string   $query Natural language query.
	 * @param string[] $types Source types to consider (template, component, shortcode, doc). Empty = all.
	 * @param int      $k     Max results.
	 * @return array<int,array{source_type:string,source_id:string,source_label:string,content:string,score:float}>
	 */
	public static function retrieve( $query, array $types = array(), $k = 6 ) {
		$t_start = microtime( true );
		$query   = trim( (string) $query );
		self::$last_diagnostic = array(); // reset

		if ( '' === $query ) {
			error_log( '[AMB-RAG] retrieve: SKIP — empty query' );
			self::record_diagnostic( array( 'skip_reason' => 'empty_query' ) );
			return array();
		}
		if ( ! AMBLD_Rag_Indexer::can_embed() ) {
			error_log( '[AMB-RAG] retrieve: SKIP — embedder unavailable (parent Antimanual provider not configured)' );
			self::record_diagnostic( array(
				'skip_reason' => 'no_embedder',
				'query_chars' => strlen( $query ),
			) );
			return array();
		}
		$qvec_raw = \Antimanual\AIProvider::generate_embedding( $query );
		if ( is_wp_error( $qvec_raw ) ) {
			error_log( '[AMB-RAG] retrieve: SKIP — query embed WP_Error: ' . $qvec_raw->get_error_message() );
			self::record_diagnostic( array(
				'skip_reason' => 'embed_error',
				'query_chars' => strlen( $query ),
				'error'       => $qvec_raw->get_error_message(),
			) );
			return array();
		}
		if ( empty( $qvec_raw ) || ! is_array( $qvec_raw ) ) {
			error_log( '[AMB-RAG] retrieve: SKIP — query embed returned empty/invalid' );
			self::record_diagnostic( array(
				'skip_reason' => 'embed_empty',
				'query_chars' => strlen( $query ),
			) );
			return array();
		}
		$qvec  = array_map( 'floatval', $qvec_raw );
		$qnorm = self::norm( $qvec );
		if ( $qnorm <= 0 ) {
			error_log( '[AMB-RAG] retrieve: SKIP — query vector has zero norm' );
			self::record_diagnostic( array(
				'skip_reason' => 'zero_norm',
				'query_chars' => strlen( $query ),
				'query_dim'   => count( $qvec ),
			) );
			return array();
		}

		$rows = AMBLD_Rag_Store::load_all( $types );
		if ( empty( $rows ) ) {
			error_log( sprintf(
				'[AMB-RAG] retrieve: NO ROWS — store empty for types=[%s]. Seed bootstrap may not have run yet; see /preferences/rag-status.',
				implode( ',', $types )
			) );
			self::record_diagnostic( array(
				'skip_reason' => 'no_rows',
				'query_chars' => strlen( $query ),
				'query_dim'   => count( $qvec ),
			) );
			return array();
		}

		$query_dim = count( $qvec );
		$query_preview = substr( preg_replace( '/\s+/', ' ', $query ), 0, 80 );
		error_log( sprintf(
			'[AMB-RAG] retrieve: query="%s" | query_dim=%d | types=[%s] | rows_loaded=%d',
			$query_preview,
			$query_dim,
			implode( ',', $types ),
			count( $rows )
		) );

		$scored = array();
		foreach ( $rows as $row ) {
			$rvec = $row['embedding'];
			if ( count( $rvec ) !== count( $qvec ) ) {
				// Embedding dimension mismatch — likely a model change since
				// the row was indexed. Skip rather than score garbage.
				continue;
			}
			$rnorm = self::norm( $rvec );
			if ( $rnorm <= 0 ) {
				continue;
			}
			$dot = 0.0;
			foreach ( $qvec as $i => $v ) {
				$dot += $v * $rvec[ $i ];
			}
			$cosine = $dot / ( $qnorm * $rnorm );

			// Helpfulness rerank — small, bounded, reversible. Multiplier
			// stays inside [0.75, 1.5] regardless of helpfulness column
			// extremes; cosine still dominates and this only matters as a
			// tiebreaker for closely-ranked chunks. Range chosen so a fresh
			// chunk (helpfulness = 0) is unchanged.
			$helpfulness = isset( $row['helpfulness'] ) ? (int) $row['helpfulness'] : 0;
			$multiplier  = 1.0 + max( -0.25, min( 0.5, $helpfulness / 20 ) );
			$score       = $cosine * $multiplier;

			$scored[] = array(
				'id'           => isset( $row['id'] ) ? (int) $row['id'] : 0,
				'source_type'  => $row['source_type'],
				'source_id'    => $row['source_id'],
				'source_label' => $row['source_label'],
				'content'      => $row['content'],
				'score'        => $score,
				'cosine'       => $cosine,
				'helpfulness'  => $helpfulness,
			);
		}

		usort(
			$scored,
			static function ( $a, $b ) {
				if ( $a['score'] === $b['score'] ) {
					return 0;
				}
				return $a['score'] < $b['score'] ? 1 : -1;
			}
		);

		$top = array_slice( $scored, 0, max( 1, (int) $k ) );

		// Mark the returned rows as recently used. Cheap UPDATE, fire-and-forget.
		if ( ! empty( $top ) ) {
			AMBLD_Rag_Store::touch_used( array_map(
				static function ( $hit ) {
					return isset( $hit['id'] ) ? (int) $hit['id'] : 0;
				},
				$top
			) );
		}

		// Emit a single-line summary of what we returned. Greppable in PHP
		// error log via `grep AMB-RAG /path/to/wp-content/debug.log`.
		$scored_count = count( $scored );
		$skipped_dim  = count( $rows ) - $scored_count;
		$summary      = array();
		foreach ( $top as $hit ) {
			$summary[] = sprintf(
				'%s/%s@%.3f',
				(string) $hit['source_type'],
				(string) $hit['source_id'],
				(float) $hit['score']
			);
		}
		$elapsed_ms = (int) round( ( microtime( true ) - $t_start ) * 1000 );
		error_log( sprintf(
			'[AMB-RAG] retrieve: scored=%d/%d (skipped %d dim-mismatch) | top_%d=[%s] | %dms',
			$scored_count,
			count( $rows ),
			$skipped_dim,
			count( $top ),
			implode( ' ', $summary ),
			$elapsed_ms
		) );

		// Diagnostic record for the harness / SSE consumer. Reachable via
		// `AMBLD_Rag_Retriever::last_diagnostic()`. `skip_reason` stays empty
		// on the success path, even when `hits` is 0 (e.g. every row was
		// dim-mismatched — that's a successful-but-fruitless call, not a
		// skip). The frontend distinguishes "skipped" from "scored 0" by
		// reading skip_reason directly.
		self::record_diagnostic( array(
			'skip_reason' => '',
			'query_chars' => strlen( $query ),
			'query_dim'   => count( $qvec ),
			'rows_loaded' => count( $rows ),
			'scored'      => $scored_count,
			'skipped_dim' => $skipped_dim,
			'hits'        => count( $top ),
			'elapsed_ms'  => $elapsed_ms,
		) );

		return $top;
	}

	private static function norm( array $vector ) {
		$sum = 0.0;
		foreach ( $vector as $v ) {
			$sum += $v * $v;
		}
		return sqrt( $sum );
	}

	/**
	 * Render retrieved chunks into a system-prompt fragment.
	 */
	public static function format_for_prompt( array $hits ) {
		if ( empty( $hits ) ) {
			return '';
		}
		$lines = array( "## Retrieved examples (top matches from this site's templates/components/docs):" );
		foreach ( $hits as $hit ) {
			$label = $hit['source_label'] ?: ( $hit['source_type'] . ':' . $hit['source_id'] );
			$lines[] = sprintf( '— [%s] %s', strtoupper( $hit['source_type'] ), $label );
			$lines[] = trim( $hit['content'] );
			$lines[] = '';
		}
		return implode( "\n", $lines );
	}

	/**
	 * Compact summary suitable for the SSE `sources_used` event.
	 */
	public static function summarize_sources( array $hits ) {
		$out = array();
		foreach ( $hits as $hit ) {
			$out[] = array(
				'type'  => (string) $hit['source_type'],
				'id'    => (string) $hit['source_id'],
				'label' => (string) ( $hit['source_label'] ?: $hit['source_id'] ),
				'score' => round( (float) $hit['score'], 4 ),
			);
		}
		return $out;
	}
}
