<?php

declare(strict_types=1);

namespace WordPress\CloudflareWorkersAiProvider\Embeddings;

use WordPress\CloudflareWorkersAiProvider\Gateway;
use WordPress\CloudflareWorkersAiProvider\Usage\UsageTracker;

use function WordPress\CloudflareWorkersAiProvider\get_account_id;
use function WordPress\CloudflareWorkersAiProvider\get_api_key;
use function WordPress\CloudflareWorkersAiProvider\get_plugin_settings;

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

/**
 * Finds the parts of posts that answer a question (2.0: better answers).
 *
 * Meaning alone misses loose questions: "something for dinner" scores every recipe about the
 * same, so the recipes that say "dinner" never surface. So two lists are made and merged:
 *
 * 1. by meaning: the content index's closest parts;
 * 2. by words: posts whose title and text contain the question's own words (common words
 *    left out, rarer words counting more), each with its part that has most of them.
 *
 * They are merged by reciprocal rank fusion (a part high in either list rises). Optionally a
 * reranking model on Workers AI then orders the best 20 by how well each answers the question.
 */
final class Retriever
{
    public const MODE_VECTOR = 'vector';
    public const MODE_HYBRID = 'hybrid';
    public const MODE_RERANK = 'rerank';

    public const RERANK_MODEL = '@cf/baai/bge-reranker-base';

    /**
     * The reranked score a part needs by default. On ragdemo (600 posts, 28 questions) every
     * right answer scored 0.11 or more, off-topic questions 0.05 or less.
     */
    public const RERANK_MIN = 0.08;

    /** Price per million input tokens, for the neuron estimate when Workers AI reports none. */
    private const RERANK_USD_PER_MILLION = 0.0031;

    /** Reciprocal rank fusion constant: the usual 60. */
    private const RRF_K = 60;

    /** Words that say little about what is asked. */
    private const STOPWORDS = ['the', 'and', 'for', 'are', 'but', 'not', 'you', 'all', 'any', 'can', 'had', 'her', 'was', 'one', 'our', 'out', 'has', 'him', 'his', 'how', 'its', 'may', 'new', 'now', 'see', 'two', 'who', 'did', 'get', 'got', 'let', 'say', 'she', 'too', 'use', 'way', 'what', 'when', 'where', 'which', 'with', 'this', 'that', 'there', 'their', 'them', 'they', 'then', 'than', 'these', 'those', 'have', 'from', 'into', 'about', 'some', 'something', 'anything', 'would', 'could', 'should', 'will', 'your', 'yours', 'make', 'made', 'like', 'just', 'also', 'more', 'most', 'much', 'many', 'very', 'want', 'need', 'give', 'show', 'tell', 'find', 'know', 'does', 'doing', 'good', 'best', 'please', 'site', 'post', 'posts', 'any', 'me', 'my', 'is', 'it', 'in', 'on', 'of', 'to', 'a', 'an', 'or', 'do', 'be', 'i'];

    /**
     * The best parts for a question.
     *
     * @return list<array{id: int, chunk: int, title: string, url: string, text: string, score: float, vector: float, words: float, rerank: float|null}>|\WP_Error
     */
    public static function find(string $question, int $limit, string $task, string $mode = self::MODE_HYBRID)
    {
        $question = trim($question);
        $byKey    = [];

        // 1. By meaning.
        $vector = Ingest::search($question, max(16, $limit * 3), false, $task);
        if (is_wp_error($vector)) {
            return $vector;
        }
        $rank = 0;
        foreach ($vector as $m) {
            if (strpos((string) $m['id'], Extractor::siteKey() . ':p:') !== 0) {
                continue; // posts and pages only, not images
            }
            $chunk = (int) substr((string) $m['id'], (int) strrpos((string) $m['id'], ':') + 1);
            $key   = (int) $m['object_id'] . ':' . $chunk;
            $byKey[$key] = self::part((int) $m['object_id'], $chunk, (string) $m['text'], (string) $m['title'], (string) $m['url']);
            $byKey[$key]['vector'] = (float) $m['score'];
            $byKey[$key]['score'] += 1 / (self::RRF_K + (++$rank));
        }

        // 2. By words.
        if ($mode !== self::MODE_VECTOR) {
            $rank = 0;
            foreach (self::byWords($question, max(10, $limit * 2)) as [$postId, $weight]) {
                [$chunk, $text] = self::closestPart($postId, $question);
                if ($text === '') {
                    continue;
                }
                $key = $postId . ':' . $chunk;
                if (!isset($byKey[$key])) {
                    $byKey[$key] = self::part($postId, $chunk, $text, '', '');
                }
                $byKey[$key]['words'] = $weight;
                $byKey[$key]['score'] += 1 / (self::RRF_K + (++$rank));
            }
        }

        $parts = array_values($byKey);
        usort($parts, static fn(array $a, array $b): int => $b['score'] <=> $a['score']);
        $parts = array_slice($parts, 0, max(20, $limit));

        // 3. Reranked.
        if ($mode === self::MODE_RERANK && $parts !== []) {
            $scores = self::rerank($question, array_column($parts, 'text'), $task);
            if (is_array($scores)) {
                foreach ($parts as $i => $p) {
                    $parts[$i]['rerank'] = $scores[$i] ?? 0.0;
                }
                usort($parts, static fn(array $a, array $b): int => ($b['rerank'] ?? 0) <=> ($a['rerank'] ?? 0));
            }
        }

        foreach ($parts as $i => $p) {
            if ($p['title'] === '') {
                $parts[$i]['title'] = html_entity_decode(get_the_title($p['id']), ENT_QUOTES, 'UTF-8');
                $parts[$i]['url']   = (string) get_permalink($p['id']);
            }
        }
        return array_slice($parts, 0, $limit);
    }

    /**
     * @return array{id: int, chunk: int, title: string, url: string, text: string, score: float, vector: float, words: float, rerank: float|null}
     */
    private static function part(int $id, int $chunk, string $text, string $title, string $url): array
    {
        return ['id' => $id, 'chunk' => $chunk, 'title' => $title, 'url' => $url, 'text' => $text, 'score' => 0.0, 'vector' => 0.0, 'words' => 0.0, 'rerank' => null];
    }

    /**
     * The question's words that carry meaning: lower case, three letters or more (or a
     * number), common words left out, at most eight.
     *
     * @return list<string>
     */
    public static function words(string $question): array
    {
        $words = preg_split('/[^\p{L}\p{N}]+/u', mb_strtolower($question)) ?: [];
        $words = array_filter($words, static fn(string $w): bool => $w !== '' && (mb_strlen($w) >= 3 || ctype_digit($w)) && !in_array($w, self::STOPWORDS, true));
        return array_slice(array_values(array_unique($words)), 0, 8);
    }

    /**
     * Posts whose title or text contain the question's words, best first: each word counts by
     * how rare it is among the matching posts, a title hit three times a text hit, and many hits
     * in a long text less than a few in a short one.
     *
     * @return list<array{0: int, 1: float}> post ID => weight
     */
    public static function byWords(string $question, int $limit): array
    {
        global $wpdb;
        $words = self::words($question);
        if ($words === []) {
            return [];
        }
        $types = array_values(array_filter((array) (get_plugin_settings()['embed_post_types'] ?? ['post', 'page']), 'post_type_exists'));
        if ($types === []) {
            return [];
        }
        // One fixed clause per word, made only of placeholders; the words go in through prepare().
        $clauses = [];
        $args    = [];
        foreach ($words as $w) {
            $clauses[] = '(post_title LIKE %s OR post_content LIKE %s)';
            $esc    = '%' . $wpdb->esc_like($w) . '%';
            array_push($args, $esc, $esc);
        }
        $typeIn = implode(',', array_fill(0, count($types), '%s'));
        $where   = implode(' OR ', $clauses);
        // phpcs:ignore WordPress.DB.DirectDatabaseQuery, WordPress.DB.PreparedSQL.InterpolatedNotPrepared, WordPress.DB.PreparedSQL.NotPrepared, WordPress.DB.PreparedSQLPlaceholders.UnfinishedPrepare, PluginCheck.Security.DirectDB.UnescapedDBParameter -- $typeIn and $where are fixed %s placeholders built above; every value goes through prepare()
        $rows = $wpdb->get_results($wpdb->prepare("SELECT ID, post_title, post_content FROM {$wpdb->posts} WHERE post_status = 'publish' AND post_password = '' AND post_type IN ({$typeIn}) AND ({$where}) ORDER BY ID DESC LIMIT 400", ...array_merge($types, $args)), ARRAY_A);
        if (!is_array($rows) || $rows === []) {
            return [];
        }

        $docs = [];
        $df   = array_fill_keys($words, 0);
        foreach ($rows as $r) {
            $title = mb_strtolower(wp_strip_all_tags((string) $r['post_title']));
            $text  = mb_strtolower(wp_strip_all_tags((string) $r['post_content']));
            $hits  = [];
            foreach ($words as $w) {
                $t = substr_count($title, $w);
                $c = substr_count($text, $w);
                if ($t + $c > 0) {
                    $df[$w]++;
                    $hits[$w] = [$t, $c];
                }
            }
            $docs[(int) $r['ID']] = [$hits, max(50, str_word_count($text))];
        }
        $n      = count($docs);
        $scores = [];
        foreach ($docs as $id => [$hits, $length]) {
            $score = 0.0;
            foreach ($hits as $w => [$t, $c]) {
                $idf  = log(1 + $n / max(1, $df[$w]));
                $tf   = 3 * $t + $c;
                $norm = $tf / ($tf + 1.2 * (0.25 + 0.75 * $length / 400)); // BM25-like saturation
                $score += $idf * $norm;
            }
            // Posts that hold more of the question's words first.
            $scores[$id] = $score * (count($hits) / count($words));
        }
        arsort($scores);
        $out = [];
        foreach (array_slice($scores, 0, $limit, true) as $id => $w) {
            if (Selection::isEligible(get_post((int) $id))) {
                $out[] = [(int) $id, round((float) $w, 4)];
            }
        }
        return $out;
    }

    /**
     * A post's part that shares most words with the question (the first part when none do).
     *
     * @return array{0: int, 1: string} [part number, its text]
     */
    public static function closestPart(int $postId, string $question): array
    {
        $post = get_post($postId);
        if (!$post instanceof \WP_Post) {
            return [0, ''];
        }
        $doc    = Extractor::extract($post);
        $chunks = Chunker::chunk($doc['title'], $doc['body'], (int) (Models::get(Models::current())['max_tokens'] ?? 512));
        if ($chunks === []) {
            return [0, ''];
        }
        $words = self::words($question);
        $best  = 0;
        $most  = 0;
        foreach ($chunks as $i => $chunk) {
            $text = mb_strtolower((string) $chunk['text']);
            $hits = 0;
            foreach ($words as $w) {
                $hits += substr_count($text, $w);
            }
            if ($hits > $most) {
                $best = (int) $i;
                $most = $hits;
            }
        }
        return [$best, (string) $chunks[$best]['text']];
    }

    /**
     * How well each text answers the question, from the reranking model (0 to 1), in the
     * texts' order; null when it could not be asked (the order then stays as merged).
     *
     * @param list<string> $texts
     * @return list<float>|null
     */
    public static function rerank(string $question, array $texts, string $task): ?array
    {
        if (get_account_id() === '' || get_api_key() === '' || $texts === []) {
            return null;
        }
        $contexts = array_map(static fn(string $t): array => ['text' => mb_substr($t, 0, 2000)], $texts);
        $url      = Gateway::runUrl(self::RERANK_MODEL);
        $headers  = ['Authorization' => 'Bearer ' . get_api_key(), 'Content-Type' => 'application/json'];
        if (Gateway::isGatewayUrl($url)) {
            $headers += Gateway::headers($task, false, self::RERANK_MODEL);
        }
        $r = wp_remote_post($url, ['timeout' => 15, 'headers' => $headers, 'body' => (string) wp_json_encode(['query' => mb_substr($question, 0, 1000), 'contexts' => $contexts, 'top_k' => count($contexts)])]);
        if (is_wp_error($r) || (int) wp_remote_retrieve_response_code($r) !== 200) {
            return null;
        }
        $data   = json_decode((string) wp_remote_retrieve_body($r), true);
        $result = is_array($data) ? ($data['result'] ?? []) : [];
        $list   = is_array($result['response'] ?? null) ? $result['response'] : [];
        $usage  = is_array($result['usage'] ?? null) ? $result['usage'] : [];
        $tokens = (int) ($usage['prompt_tokens'] ?? 0);
        if ($tokens === 0) {
            $tokens = (int) ceil(mb_strlen($question . implode(' ', array_column($contexts, 'text'))) / 4);
        }
        UsageTracker::record(self::RERANK_MODEL, round($tokens / 1000 * UsageTracker::neuronsPer1k(self::RERANK_USD_PER_MILLION), 4), $tokens, 0, $task);
        $scores = array_fill(0, count($texts), 0.0);
        foreach ($list as $item) {
            if (is_array($item) && isset($item['id'], $item['score'])) {
                $scores[(int) $item['id']] = (float) $item['score'];
            }
        }
        return $scores;
    }
}