<?php

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

/**
 * Vector search over the knowledge base: embeds a query, scans stored chunk
 * embeddings by cosine similarity, and assembles the winning chunks into
 * context blocks for the system prompt.
 */
class RSVL_Search
{
    private ?array $last_query_vector = null;

    public function get_last_query_vector(): ?array
    {
        return $this->last_query_vector;
    }

    /**
     * Search the knowledge base for chunks relevant to the query.
     *
     * @param string $query The user's question.
     * @return array Array of matching entries with 'content', 'source_url', 'similarity', 'chunk_index', 'total_chunks'.
     */
    public function search(string $query): array
    {
        global $wpdb;

        $embedder = new RSVL_Embedder();
        if (! $embedder->is_available()) {
            return [];
        }

        // Embed the query.
        $query_vector = $embedder->embed($query);
        if (is_wp_error($query_vector)) {
            resolver_log('Failed to embed query: ' . $query_vector->get_error_message());

            return [];
        }

        $this->last_query_vector = $query_vector;

        $table      = RSVL_Schema::chunks_table();
        $threshold  = (float) apply_filters('resolver_similarity_threshold', $embedder->get_similarity_threshold());
        $max_chunks = (int) apply_filters('resolver_max_context_chunks', resolver_get_setting('max_context_chunks'));

        // Phase 1: Scan embeddings in batches to find candidates.
        $batch_size = 250;
        $offset     = 0;
        $candidates = [];

        $embedding_model = $embedder->get_model();

        do {
            $batch = $wpdb->get_results($wpdb->prepare(
                'SELECT id, embedding FROM %i WHERE embedding IS NOT NULL AND embedding_model = %s LIMIT %d OFFSET %d',
                $table,
                $embedding_model,
                $batch_size,
                $offset
            ));

            if (empty($batch)) {
                break;
            }

            foreach ($batch as $row) {
                $doc_vector = RSVL_Embedder::unpack_vector($row->embedding);
                $similarity = RSVL_Embedder::cosine_similarity($query_vector, $doc_vector);

                if ($similarity >= $threshold) {
                    $candidates[] = [
                        'id'         => (int) $row->id,
                        'similarity' => $similarity,
                    ];
                }
            }

            $offset += $batch_size;
        } while (count($batch) === $batch_size);

        if (empty($candidates)) {
            return [];
        }

        // Sort by similarity descending.
        usort($candidates, function ($a, $b) {
            return $b['similarity'] <=> $a['similarity'];
        });

        $min_context_chars = (int) apply_filters('resolver_min_context_chars', 4000);

        // Take top N candidates.
        $top     = array_slice($candidates, 0, $max_chunks);
        $top_ids = array_column($top, 'id');

        // Phase 2: Fetch content for the initial top N.
        $rows = $this->fetch_chunk_rows_by_ids($top_ids);

        $total_chars = array_sum(array_map(function ($row) {
            return mb_strlen($row->content);
        }, $rows));

        // The top N is a floor, not a target: if those chunks are too thin
        // (e.g. a few short notes all ranking highly), keep pulling the
        // next-most-similar candidates — still only ones that already
        // cleared the similarity threshold in phase 1 — until the combined
        // content reaches the floor, or no more candidates remain. This
        // keeps a longer, still-relevant chunk from being crowded out just
        // because a handful of short ones happened to rank slightly higher.
        $extension_batch_size = 10;
        $next_index           = $max_chunks;

        while ($total_chars < $min_context_chars && $next_index < count($candidates)) {
            $batch_ids  = array_column(array_slice($candidates, $next_index, $extension_batch_size), 'id');
            $rows_by_id = [];
            foreach ($this->fetch_chunk_rows_by_ids($batch_ids) as $row) {
                $rows_by_id[ (int) $row->id ] = $row;
            }

            foreach ($batch_ids as $id) {
                if ($total_chars >= $min_context_chars) {
                    break;
                }
                if (! isset($rows_by_id[ $id ])) {
                    continue;
                }
                $rows[] = $rows_by_id[ $id ];
                $total_chars += mb_strlen($rows_by_id[ $id ]->content);
            }

            $next_index += $extension_batch_size;
        }

        $similarity_map = array_column($candidates, 'similarity', 'id');

        // Build results with similarity scores.
        $output = [];
        foreach ($rows as $row) {
            $output[] = [
                'id'           => (int) $row->id,
                'content'      => $row->content,
                'source_url'   => $row->source_url,
                'chunk_index'  => (int) $row->chunk_index,
                'total_chunks' => (int) $row->total_chunks,
                'source_id'    => (int) $row->source_id,
                'similarity'   => $similarity_map[ (int) $row->id ] ?? 0,
            ];
        }

        // Sort by similarity again since the DB queries don't preserve order.
        usort($output, function ($a, $b) {
            return $b['similarity'] <=> $a['similarity'];
        });

        return $output;
    }

    /**
     * Build context blocks from search results, reassembling chunks from the same source.
     *
     * @param array $results The search results from self::search().
     * @return array Array of ['source_url' => '...', 'content' => '...'] blocks.
     */
    public function build_context_blocks(array $results): array
    {
        // Group by source URL.
        $grouped = [];
        foreach ($results as $r) {
            $url = $r['source_url'];
            if (! isset($grouped[ $url ])) {
                $grouped[ $url ] = [];
            }
            $grouped[ $url ][] = $r;
        }

        // Sort each group by chunk_index and reassemble.
        $blocks = [];
        foreach ($grouped as $url => $chunks) {
            usort($chunks, function ($a, $b) {
                return $a['chunk_index'] <=> $b['chunk_index'];
            });

            $content_parts = [];
            foreach ($chunks as $chunk) {
                $content_parts[] = $chunk['content'];
            }

            $blocks[] = [
                'source_url' => $url,
                'content'    => implode("\n\n", $content_parts),
            ];
        }

        return $blocks;
    }

    /**
     * Fetch chunk content + metadata for a set of knowledge-table ids.
     */
    private function fetch_chunk_rows_by_ids(array $ids): array
    {
        global $wpdb;

        if (empty($ids)) {
            return [];
        }

        return $wpdb->get_results($wpdb->prepare(
            'SELECT id, content, source_url, chunk_index, total_chunks, source_id
             FROM %i
             WHERE id IN (' . implode(',', array_fill(0, count($ids), '%d')) . ')',
            array_merge([RSVL_Schema::chunks_table()], $ids)
        ));
    }
}