<?php
/**
 * Semantic index for the runtime tool catalog.
 *
 * Ranks tool catalog entries against a free-form query using embedding
 * cosine similarity, so tool_search works in any language instead of
 * depending on English substring matches. It batch-embeds entries through
 * Niwa_AI_Embedding_Client, persists vectors in a non-autoloaded option,
 * and compares them with cosine similarity.
 *
 * Vectors are truncated to 256 dimensions and re-normalized (valid for
 * text-embedding-3 Matryoshka models) to keep the cached option small.
 * Entry vectors are cached per text hash, so unchanged tool definitions
 * are never re-embedded and different role registries share cache entries.
 *
 * Every public method is non-fatal: when the embedding client, token, or
 * network is unavailable, rank() returns null and the caller is expected
 * to fall back to lexical scoring.
 *
 * @package Niwa_AI
 */

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

class Niwa_AI_Admin_Autonomy_Tool_Catalog_Index {

    /** Option holding the cached entry vectors (autoload=false). */
    const EMBEDDINGS_OPTION = 'niwa_ai_tool_catalog_embeddings';

    /** Stored vector dimensions after Matryoshka truncation. */
    const VECTOR_DIMENSIONS = 256;

    /** Bump to invalidate all cached vectors (format/model changes). */
    const SCHEMA_VERSION = '1';

    /** Hard cap on cached vectors before pruning to the current entry set. */
    const MAX_CACHED_VECTORS = 600;

    /**
     * Embedding client used for catalog and query vectors.
     *
     * @var Niwa_AI_Embedding_Client|null
     */
    private $embedding_client;

    /**
     * @param object|null $embedding_client Optional injected client (tests).
     */
    public function __construct( $embedding_client = null ) {
        if ( is_object( $embedding_client ) && method_exists( $embedding_client, 'embed_texts' ) ) {
            $this->embedding_client = $embedding_client;
        } elseif ( class_exists( 'Niwa_AI_Embedding_Client' ) ) {
            $this->embedding_client = new Niwa_AI_Embedding_Client();
        }
    }

    /**
     * Rank catalog entries against a query.
     *
     * @param array  $entries Array of {name: string, text: string} rows.
     * @param string $query   Free-form query in any language.
     *
     * @return array|null Map of entry name => similarity (0.0–1.0), or null
     *                    when semantic ranking is unavailable and the caller
     *                    must fall back to lexical scoring.
     */
    public function rank( $entries, $query ) {
        $query   = trim( (string) $query );
        $entries = $this->sanitize_entries( $entries );

        if ( ! $this->embedding_client || '' === $query || empty( $entries ) ) {
            return null;
        }

        $vectors = $this->get_entry_vectors( $entries );
        if ( null === $vectors ) {
            return null;
        }

        $query_vector = $this->prepare_vector( $this->embedding_client->embed_text( $query ) );
        if ( empty( $query_vector ) ) {
            return null;
        }

        $ranked = array();
        foreach ( $entries as $entry ) {
            $hash = md5( $entry['text'] );
            if ( ! isset( $vectors[ $hash ] ) ) {
                continue;
            }
            $ranked[ $entry['name'] ] = $this->cosine_similarity( $query_vector, $vectors[ $hash ] );
        }

        return empty( $ranked ) ? null : $ranked;
    }

    /**
     * Resolve vectors for all entries, embedding only cache misses.
     *
     * @param array $entries Sanitized entries.
     *
     * @return array|null Map of text hash => vector, or null on embed failure.
     */
    private function get_entry_vectors( $entries ) {
        $stored = get_option( self::EMBEDDINGS_OPTION, array() );
        $cache  = array();

        if (
            is_array( $stored )
            && isset( $stored['version'] )
            && self::SCHEMA_VERSION === (string) $stored['version']
            && isset( $stored['vectors'] )
            && is_array( $stored['vectors'] )
        ) {
            $cache = $stored['vectors'];
        }

        $needed  = array();
        $missing = array();
        foreach ( $entries as $entry ) {
            $hash            = md5( $entry['text'] );
            $needed[ $hash ] = true;
            if ( ! isset( $cache[ $hash ] ) || ! is_array( $cache[ $hash ] ) ) {
                $missing[ $hash ] = $entry['text'];
            }
        }

        if ( ! empty( $missing ) ) {
            $missing_hashes = array_keys( $missing );
            $embedded       = $this->embedding_client->embed_texts( array_values( $missing ) );

            if ( ! is_array( $embedded ) || count( $embedded ) !== count( $missing_hashes ) ) {
                return null;
            }

            foreach ( $missing_hashes as $index => $hash ) {
                $vector = $this->prepare_vector( isset( $embedded[ $index ] ) ? $embedded[ $index ] : array() );
                if ( empty( $vector ) ) {
                    return null;
                }
                $cache[ $hash ] = $vector;
            }

            if ( count( $cache ) > self::MAX_CACHED_VECTORS ) {
                $cache = array_intersect_key( $cache, $needed );
            }

            update_option(
                self::EMBEDDINGS_OPTION,
                array(
                    'version' => self::SCHEMA_VERSION,
                    'vectors' => $cache,
                ),
                false
            );
        }

        return $cache;
    }

    /**
     * Truncate to VECTOR_DIMENSIONS, L2-normalize, and round for storage.
     *
     * @param array $vector Raw embedding vector.
     *
     * @return array Prepared vector, or empty array when input is unusable.
     */
    private function prepare_vector( $vector ) {
        if ( ! is_array( $vector ) || count( $vector ) < self::VECTOR_DIMENSIONS ) {
            return array();
        }

        $vector = array_slice( $vector, 0, self::VECTOR_DIMENSIONS );
        $norm   = 0.0;
        foreach ( $vector as $value ) {
            $value = (float) $value;
            $norm += $value * $value;
        }
        $norm = sqrt( $norm );
        if ( $norm <= 0.0 ) {
            return array();
        }

        $prepared = array();
        foreach ( $vector as $value ) {
            $prepared[] = round( ( (float) $value ) / $norm, 6 );
        }

        return $prepared;
    }

    /**
     * Cosine similarity between two vectors.
     *
     * Vectors are stored normalized, so this is effectively a dot product,
     * but the full formula keeps the method safe for arbitrary input.
     *
     * @param array $a Left vector.
     * @param array $b Right vector.
     *
     * @return float Similarity in the 0.0–1.0 range (negative clamped to 0).
     */
    private function cosine_similarity( $a, $b ) {
        $len = min( count( $a ), count( $b ) );
        if ( $len <= 0 ) {
            return 0.0;
        }

        $dot    = 0.0;
        $norm_a = 0.0;
        $norm_b = 0.0;
        for ( $i = 0; $i < $len; $i++ ) {
            $va      = (float) $a[ $i ];
            $vb      = (float) $b[ $i ];
            $dot    += $va * $vb;
            $norm_a += $va * $va;
            $norm_b += $vb * $vb;
        }

        $denom = sqrt( $norm_a ) * sqrt( $norm_b );
        if ( $denom <= 0.0 ) {
            return 0.0;
        }

        return max( 0.0, $dot / $denom );
    }

    /**
     * Keep only well-formed entries with non-empty name and text.
     *
     * @param mixed $entries Raw entries.
     *
     * @return array Sanitized entries.
     */
    private function sanitize_entries( $entries ) {
        if ( ! is_array( $entries ) ) {
            return array();
        }

        $clean = array();
        foreach ( $entries as $entry ) {
            if ( ! is_array( $entry ) ) {
                continue;
            }
            $name = isset( $entry['name'] ) ? trim( (string) $entry['name'] ) : '';
            $text = isset( $entry['text'] ) ? trim( (string) $entry['text'] ) : '';
            if ( '' === $name || '' === $text ) {
                continue;
            }
            $clean[] = array(
                'name' => $name,
                'text' => $text,
            );
        }

        return $clean;
    }
}
