<?php
// Expects to be in /root/includes folder

require_once __DIR__ . '/twx_util.php';
require_once __DIR__ . '/openai_include.php';

class TWXChatKnowledgeBaseUtil
{
	// Util class with static methods relating to local Knowledge Base

    public static function getAllKnowledgeBaseEntries() {

        global $wpdb;
        $sql = "SELECT * FROM {$wpdb->prefix}twxchat_vectors";
        $query = $wpdb->prepare($sql);
        $results = $wpdb->get_results($query);
        return self::getEntriesFromResults($results);
    }

    public static function getEntriesBySourceIds($source_ids) {
        global $wpdb;

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

        // Dynamically generate placeholders for the IN clause
        $placeholders = implode(',', array_fill(0, count($source_ids), '%d'));
        $sql = "SELECT * FROM {$wpdb->prefix}twxchat_vectors WHERE id IN ($placeholders) ORDER BY id";

        $query = $wpdb->prepare($sql, ...$source_ids);
        $results = $wpdb->get_results($query);

        return self::getEntriesFromResults($results);
    }

    public static function isFeedbackSource($source) {
         return str_ends_with($source, '_feedback');
    }

    public static function getSourcesBySourceIds($source_ids) {
        global $wpdb;

        // Sanitize and validate the source IDs
        $source_ids = array_map('intval', (array) $source_ids);

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

        // Create the placeholders for the IN clause
        $placeholders = implode(',', array_fill(0, count($source_ids), '%d'));
        $sql = "SELECT * FROM {$wpdb->prefix}twxchat_vectors WHERE id IN ($placeholders) ORDER BY id";

        // Prepare the query securely
        $query = $wpdb->prepare($sql, ...$source_ids);
        $results = $wpdb->get_results($query);
        return self::getSourcesFromResults($results);
    }

    public static function getEntriesBySource($sourcesString, $strict_match = true) {
        global $wpdb;
    
        // Convert comma-separated string to array
        $sources = array_map('trim', explode(',', $sourcesString));
    
        // Define the table name with the prefix
        $table_name = $wpdb->prefix . 'twxchat_vectors';
    
        // Initialize variables
        $sql = '';
        $params = [];
    
        if ($strict_match) {
            // Strict match: use IN clause
            $placeholders = implode(',', array_fill(0, count($sources), '%s'));
            $sql = "SELECT * FROM {$table_name} WHERE source IN ($placeholders)"; // only place holders exist here, no actual values
            $params = $sources;
        } else {
            // Partial match: use multiple LIKE conditions
            $likesPlaceholders = [];
            foreach ($sources as $src) {
                $likesPlaceholders[] = "source LIKE %s";
                $params[] = '%' . $wpdb->esc_like($src) . '%';
            }
            $sql = "SELECT * FROM {$table_name} WHERE " . implode(' OR ', $likesPlaceholders); // only place holders exist here, no actual values
        }
    
        // Prepare the SQL statement
        $query = $wpdb->prepare($sql, ...$params);

        // Execute the query and get the results
        $results = $wpdb->get_results($query);
        return self::getEntriesFromResults($results);
    }
    
    public static function getEntriesNoSource() {

        // used when no Knowledge Base source is used in shortcode, ie: search whole KB
        
        global $wpdb;
        $sql = "SELECT * FROM {$wpdb->prefix}twxchat_vectors";

        // Execute the query and get the results - no preparation needed
        $results = $wpdb->get_results($sql);
        return self::getEntriesFromResults($results);
    }

    public static function getFilteredEntries($filters, $sortColName = null, $sortAscDesc = true) {

        global $wpdb;

        $wherePlaceholders = [];
        $queryParams = [];

        $columns = self::getAdminColumns();

        foreach ($filters as $colName => $colFilter) {
            if (empty($colFilter)) continue;

            // Only allow whitelisted column names
            if (!isset($columns[$colName]) || (!empty($columns[$colName]['non_db_column']))) {
                continue;
            }

            $values = array_map('trim', explode(',', $colFilter));
            $orPlaceholders = [];

            foreach ($values as $val) {
                $orPlaceholders[] = "$colName LIKE %s";
                $queryParams[] = '%' . $wpdb->esc_like($val) . '%';
            }

            if (!empty($orPlaceholders)) {
                $wherePlaceholders[] = '(' . implode(' OR ', $orPlaceholders) . ')';
            }
        }

        $table = $wpdb->prefix . 'twxchat_vectors';
        $sql = "SELECT * FROM $table";

        if (!empty($wherePlaceholders)) {
            $sql .= " WHERE " . implode(" AND ", $wherePlaceholders); // no values are in $wherePlaceholders only place holders which will be parsed by $wpdb->prepare below
        }

        // Safely append ORDER BY, ie: Only allow whitelisted column names
        if (!empty($sortColName) && isset($columns[$sortColName])) {
            if (empty($columns[$sortColName]['non_db_column'])) {
                $orderDir = $sortAscDesc ? 'ASC' : 'DESC';
                $sql .= " ORDER BY `$sortColName` $orderDir"; // $sortColName and $orderDir are derived from trusted code not user input
            }
        }


        // Prepare the final SQL with parameters
        $sql_prepared = $wpdb->prepare($sql, ...$queryParams);

        $results = $wpdb->get_results($sql_prepared); // sql is prepared if paramaters used
        return self::getEntriesFromResults($results);
    }
   
    public static function getEntriesFromResults($results) {

        if(empty($results)) return [];
        $ret = [];

        foreach ($results as $result) {
            $ret[$result->id] = [
                'embedding'  => json_decode($result->vector),
                'orig_text'  => $result->orig_text,
                'title'      => $result->title,
                'source'     => $result->source,
                'created_at' => $result->created_at,
            ];
        }
        return $ret;
    }

    public static function getSourcesFromResults($results) {

        if(empty($results)) return [];
        $ret = [];

        foreach ($results as $result) {
            $ret[] = [
                'source_id' => $result->id,
                'title'     => $result->title,
                'source'    => $result->source,
                'orig_text' => $result->orig_text,
                'created_at'=> $result->created_at,
            ];
        }
        return $ret;
    }

    public static function updateEntryById($apiKey, $kbId, $data){

        global $wpdb;

        $table_name = $wpdb->prefix . 'twxchat_vectors';

        $newData = [];  
        $newData['title'] = $data['title'];
        $newData['source'] = $data['source'];
        $newData['orig_text'] = $data['orig_text'];
        $newData['vector'] = json_encode(TWXChatUtil::getEmbedding($apiKey, $data['orig_text']));
        $newData['created_at'] = current_time('mysql');
        $where = ['id' => $kbId];

        $updated = $wpdb->update($table_name, $newData, $where);

        if ($updated !== false) {
            //TWXChatUtil::log('Row updated successfully.');
        } else {
            TWXChatUtil::adminLog('TWXChatKnowledgeBaseUtil::updateEntryById() Failed to update row.');
        }

    }

    public static function hasOrigTextChanged($source, $text){

        $entries = self::getEntriesBySource($source, true);
        if(empty($entries)) return true; // new source
        //$entry = $entries[0];
        $entry = reset($entries); // get first entry
        return ($entry['orig_text'] ?? '') !== $text;
    }

    public static function storeEntry($apiKey, $source, $title, $text, $replace = true, $old_source = '') {
	
        try {
            // First, delete any existing Knowledge Base entry(s) with the same source. $old_source can be used to remove old equivalent sources eg: those only with url and no title
            if($replace) {
				if(!empty($old_source)) self::deleteEntryBySource($old_source);
				self::deleteEntryBySource($source);
			}
    
            // Generate the entry using the API key and text
            $embedding = TWXChatUtil::getEmbedding($apiKey, $text);
    
            // Get the global WordPress database object
            global $wpdb;
    
            // Define the table name
            $table_name = $wpdb->prefix . 'twxchat_vectors';
    
            // Prepare data to be inserted
            $data = [
                'title' => $title,
                'source' => $source,
                'orig_text' => $text,
                'vector' => json_encode($embedding),
                'created_at' => current_time('mysql'),
            ];
    
            // Insert the new entry into the database
            $result = $wpdb->insert($table_name, $data);
    
            if ($result) {
                //TWXChatUtil::log("Knowledge Base entry stored for $source");
            } else {
                TWXChatUtil::adminLog("TWXChatKnowledgeBaseUtil::storeEntry() Failed to store Knowledge Base entry for $source");
            }
        } catch (Exception $ex) {
            $error = "Error storing Knowledge Base entry for $source: " . $ex->getMessage();
            TWXChatUtil::adminLog("TWXChatKnowledgeBaseUtil::storeEntry() $error",'a',true);
            throw new Exception(esc_html($error));
        }
        
    }
    
    public static function deleteAllKnowledgeBaseEntries($including_feedback_items = false) {
        TWXChatUtil::adminLog('TWXChatKnowledgeBaseUtil::deleteAllKnowledgeBaseEntries()');
    
        global $wpdb;
    
        // Define the table name (with WordPress prefix)
        $table_name = $wpdb->prefix . 'twxchat_vectors';
    
        // Prepare the delete query and parameters
        if ($including_feedback_items) {
            // Delete all entries, no parameters needed
            $query = "DELETE FROM {$table_name}";
            $prepared_query = $query;
        } else {
            // Delete entries where source NOT LIKE '%_feedback'
            $like_pattern = '%_feedback';
            $query = "DELETE FROM {$table_name} WHERE source NOT LIKE %s";
            $prepared_query = $wpdb->prepare($query, $like_pattern);
        }
    
        // Execute the query
        $result = $wpdb->query($prepared_query);
    
        // Return the number of affected rows
        return $result;
    }
        
    public static function deleteEntryBySource($sourceValue) {
        global $wpdb;
    
        // Define the table name (with WordPress prefix)
        $table_name = $wpdb->prefix . 'twxchat_vectors';
    
        // Prepare the delete query safely with placeholder
        $query = $wpdb->prepare(
            "DELETE FROM {$table_name} WHERE source = %s",
            $sourceValue
        );
    
        try {
            // Execute the query
            $result = $wpdb->query($query);
    
            // Return the number of affected rows (or 0 if none)
            return $result;
        } catch (Exception $ex) {
            // Log the error and return false
            TWXChatUtil::adminLog("deleteAllKnowledgeBaseEntries::deleteEntryBySource() error: " . $ex->getMessage(), 'a', true);
            return false;
        }
    }
    
    public static function deleteEntryById($kbId) {
        global $wpdb;
    
        $table_name = $wpdb->prefix . 'twxchat_vectors';
    
        // Ensure $kbId is an integer for safety
        $kbId = (int) $kbId;
    
        // Prepare the delete query using %d for integers
        $query = $wpdb->prepare("DELETE FROM {$table_name} WHERE id = %d", $kbId);
    
        try {
            $result = $wpdb->query($query);
            // Return the number of affected rows
            return $result;
        } catch (Exception $ex) {
            TWXChatUtil::adminLog("deleteAllKnowledgeBaseEntries::deleteEntryById() error: " . $ex->getMessage(), 'a', true);
            return false;
        }
    }
    
	public static function searchSimilarKnowledgeBaseEntries(&$tiktokenProvider, $apiKey, $userMessages, $entries, $includeEntries = [], $model = null, $maxTokens = -1, $minSimilarity = 0.8, $topN = 4) 	{

		/*
		Compare message/messages to Knowledge Base entries and return similar ones (including their original text and source) for use in LLM prompts.
        $userMessages can be single string or array of strings representing user messages in reverse order, ie: most recent first etc
        $KnowledgeBaseKeep keyed array contains already identified similar KnowledgeBase entries we want to return in the results (eg: from previous sources for follow-up q's)
		
		$entries should be an array as follows eg: [
			['id1' => ['source' => '{source eg: url/file name etc}', 'orig_text' => '{original content eg: markdown/html}', 'embedding' => '{embedding of orig_text}'] ],
			['id2' => ['source' => '{source eg: url/file name etc}', 'orig_text' => '{original content eg: markdown/html}', 'embedding' => '{embedding of orig_text}'] ],
			...
		]
		If $maxTokens and $model are passed in, top similar entries will be gathered until max tokens for the specified $model is reached and then returned otherwise topN similar entries will be returned
		Only entries of similarity of $minSimilarity or greater will be returned (decimal values between 0 and 1)
		
		returns: similar entries with the same structure as $entries input but with a similarity_score also
		*/

        //TWXChatUtil::log("KB.searchSimilarKnowledgeBaseEntries:: minSimilarity: $minSimilarity");

        //$isFollowUpQuery = TWXChatUtil::isFollowUpQuery($currentUserMsg);

        if(!is_array($userMessages)){
            $userMessages = [$userMessages];
        }

        $similarities = [];
        $text = "";
        foreach($userMessages as $message){
            $text = trim("$message $text"); // prepend current message to text
            try {
                $textEmbedding = TWXChatUtil::getEmbedding($apiKey, $text);
            } catch (Exception $ex){
                throw new Exception('Error getting OpenAI embedding: ' . $ex->getMessage(), 0, $ex);
            }
    
            foreach($entries as $id => $entry){
                if(empty($entry['embedding'])){
                    TWXChatUtil::log('TWXChatKnowledgeBaseUtil::searchSimilarKnowledgeBaseEntries() empty embedding for kbId: '.$id.' '.$entry['orig_text']);
                    continue;
                }
                $similarity = TWXChatUtil::cosineSimilarity($textEmbedding, $entry['embedding']);
                if($similarity >= $minSimilarity || isset($includeEntries[$id])){
                    if(isset($includeEntries[$id])){
                        $entry['similarity_score'] = 0.99; // ensure this gets included and is given priority
                    }else{
                        $entry['similarity_score'] = $similarity;
                    }
                    $similarities[$id] = $entry;
                }
            }

            if(!empty($similarities)) break; // otherwise keep going with messages until we get similarities or end messages
        }

        if(empty($similarities)) return [];

		// reverse sort on similarity_score keeping keys in tact
		uasort($similarities, function ($a, $b) {
			if(isset($a['levenshtein_similarity']) && isset($b['levenshtein_similarity'])){
				$maxA = max($a['levenshtein_similarity'], $a['similarity_score']);
				$maxB = max($b['levenshtein_similarity'], $b['similarity_score']);
				return $maxB <=> $maxA;
			}else{
				return $b['similarity_score'] <=> $a['similarity_score'];
			}
		});

        $ret = [];
		if($maxTokens>0 && !empty($model)){
			// get until max tokens
			$tokenCount = 0; $top_similarities = [];
			foreach($similarities as $id => $similarity){
				$numTokens = count(TWXChatUtil::getTokens($tiktokenProvider, $model, $similarity['orig_text']));
				if($tokenCount + $numTokens > $maxTokens) break;
				$tokenCount += $numTokens;
				$similarity['est_tokens'] = $numTokens;
				$top_similarities[$id] = $similarity;
			}
			$ret = $top_similarities;
			
		}else{
			// get topN
			$ret = array_slice($similarities, 0, $topN, true);
		}

		return $ret;
	}
	
	public static function getContextFromEntries($entries, $showSources = false){
        if(empty($entries)) return '';
        $context = '';
		foreach($entries as $id => $entry){
			if($showSources){
				$context .= "source_id: $id, content: " . $entry['orig_text'] . "\n";
			}else{
				$context .= $entry['orig_text'] . "\n";
			}
			$sources[$id] = [
				'source_id' => $id, 
				'source' => $entry['source'], 
				'similarity_score' => $entry['similarity_score'], 
				'content_length_chars' => strlen($entry['orig_text']),
				'last_modified' => $entry['created_at'],
			];
		}
        return $context;
	}

    public static function addLevenshteinSimilarityToEntries($message, &$entries){
        // may not be necessary as embeddings comparisons are probably better than word distances
        foreach($entries as &$entry){
            $entry['levenshtein_similarity'] = self::levenshteinSimilarity($message, $entry['orig_text']);
        }
    }

    public static function getAdminColumns(){
        return [
            'title' => ['display'=>'Title','sortable'=>true, 'canFilter'=>true,'canEdit'=>true],
            'source' => ['display'=>'Source','sortable'=>true, 'canFilter'=>true,'canEdit'=>true,'defaultVal'=>'_feedback','required'=>true],
            'created_at' => ['display'=>'Last Modified','sortable'=>true],
            'orig_text' => ['display'=>'Content', 'parseFunction'=>'adminTableParseContent','canFilter'=>true,'canEdit'=>true, 'similarQueryable'=>true,'required'=>true],
            'similarity' => ['display'=>'Similarity', 'parseFunction'=>'adminTableParseSimilarity', 'sortable'=>true, 'non_db_column' => true],
            'isFeedback' => ['display'=>'Is Feedback?','parseFunction'=>'adminTableParseIsFeedback', 'non_db_column' => true],
        ];
    }

    public static function adminTableParseContent($entry){
        return "<pre>{$entry['orig_text']}</pre>";
    }

    public static function adminTableParseIsFeedback($entry){
        if(!isset($entry['source'])) return 'No';
        return str_ends_with($entry['source'], 'feedback') ? 'Yes' : 'No';
    }

    public static function adminTableParseSimilarity($entry){
        if(!isset($entry['similarity'])) return '';
        return round($entry['similarity'] * 100) . '%';
    }

    public static function sortEntries($entries, $sortColName, $sortAsc){
        // If we want to perform post sort
        $columns = self::getAdminColumns();
        if(!isset($columns[$sortColName])) return $entries;
        if(isset($columns[$sortColName]['sortFunction'])){
            $functionName = $columns[$sortColName]['sortFunction'];
            if (method_exists(self::class, $functionName)) {
                $entries = self::$functionName($entries, $sortAsc);
            }
        }else{
            uasort($entries, function($a, $b) use($sortColName, $sortAsc) {
                if($sortAsc) {
                    return $a[$sortColName] <=> $b[$sortColName];
                }else{
                    return $b[$sortColName] <=> $a[$sortColName];
                }
            });
        }
        return $entries;
    }

    public static function getAdminHtml($entries, $cols, $showColumns, $filters, $sortColName, $sortAsc, $similarQueryColIndexes){

        function getClass($iCol,$showColumns,$colData, $isFilterRow = false){
            $classes = [];
            if(!empty($showColumns) && !in_array(strval($iCol), $showColumns)) {
                $classes[] = 'hidden';
            }
            if(isset($colData['sortable']) && $colData['sortable'] === true){
                $classes[] = 'sortable';
            }
            if(isset($colData['similarQueryable']) && $colData['similarQueryable'] === true && $isFilterRow){
                $classes[] = 'similar-queryable';
            }
            
            if(empty($classes)) return '';
            return ' class="'.implode(' ',$classes).'"';
        }

        function getFilterFieldHtml($iCol, $filters, $similarQueryColIndexes, $similarQueryable = false){
            $val = isset($filters[$iCol]) ? $filters[$iCol] : '';
            $similarSearch = '';
            if($similarQueryable){
                $checked = in_array($iCol ,$similarQueryColIndexes) ? ' checked' : '';
                $similarSearch = '<input type="checkbox" class="similar-query-filter" name="similarquery-filter-field-'.$iCol.'" column-index="'.$iCol.'"'.$checked.' />&nbsp;Query&nbsp;';
            }
            return $similarSearch.'<input class="filter-field" type="text" column-index="'.$iCol.'" name="filter-field-'.$iCol.'" value="'.$val.'">';
        }

        // table header row
        $html = '<table id="knowledgeBaseEntriesTable" border="1"><thead><tr><th><input type="checkbox" id="twx-chat-knowledgebase-select-chk"></th>';
        $iCol = 0;
        foreach($cols as $colName => $colData){
            $sortArrow = ($colName === $sortColName) ? '&nbsp;'.($sortAsc ? '▲' : '▼') : '';
            $html .= '<th'.getClass($iCol,$showColumns,$colData).'>'.$colData['display'].$sortArrow.'</th>';
            $iCol++;
        }
        $html .= '</tr></thead>';

        // filter fields row
        $html .= '<tbody><tr><td></td>';
        $iCol = 0;
        foreach($cols as $colName => $colData){
            if(isset($colData['canFilter']) && $colData['canFilter'] === true){
                $similarQueryable = isset($colData['similarQueryable']) && $colData['similarQueryable'];
                $html .= '<td'.getClass($iCol,$showColumns,$colData, true).'>'.getFilterFieldHtml($iCol, $filters, $similarQueryColIndexes, $similarQueryable).'</td>';
            }else{
                $html .= '<td'.getClass($iCol,$showColumns,$colData, true).'></td>';
            }
            $iCol++;
        }
        $html .= '</tr>';

        // data rows
        foreach($entries as $kbId => $entry){
            $checked = '';
            $html .= '<tr data-row-id="'.$kbId.'"><td><input type="checkbox"'.$checked.'></td>';
            $iCol = 0;

            foreach($cols as $colName => $colData){

                $html .= '<td data-colname="'.$colName.'"';

                if(isset($colData['parseFunction'])){
                    $functionName = $colData['parseFunction'];
                    if (method_exists(self::class, $functionName)) {
                        $result = self::$functionName($entry); // pass in whole entry to fn by name
                        $html .= ' '.getClass($iCol,$showColumns,$colData).'>'.$result.'</td>';
                    } else {
                        $html .= ' '.getClass($iCol,$showColumns,$colData).'>err</td>';
                    }
                }else{
                    $html .= ' '.getClass($iCol,$showColumns,$colData).'>'.$entry[$colName].'</td>';
                }
                $iCol++;
            }

            $html .= '</tr>';
        }
        $html .= '</tbody></table>';
    
        return $html;
    }
    
    private static function levenshteinSimilarity($needle, $haystack) {
        if(stripos($haystack,$needle)!== false) return 1;
        $needle = strtolower($needle);
        $haystack = strtolower($haystack);
        
        // Normalize haystack and needle
        $needle_len = strlen($needle);
        $haystack_len = strlen($haystack);
    
        // Initialize minimum distance and best match length
        $min_distance = PHP_INT_MAX;
        $best_match_len = 0;
    
        // Iterate through each possible substring in the haystack
        for ($i = 0; $i <= $haystack_len - $needle_len; $i++) {
            $substring = substr($haystack, $i, $needle_len);
            $distance = levenshtein($needle, $substring);
    
            if ($distance < $min_distance) {
                $min_distance = $distance;
                $best_match_len = strlen($substring);
            }
        }
    
        // Calculate the percentage match based on the best match length
        $max_len = max($needle_len, $best_match_len);
        $similarity = (1 - ($min_distance / $max_len));
        return $similarity;
    }

}