<?php

namespace Limb_Chatbot\Includes\Services\Knowledge;

use Limb_Chatbot\Includes\Data_Objects\Dataset_Entry;
use Limb_Chatbot\Includes\Data_Objects\Message;
use Limb_Chatbot\Includes\Utilities\Chatbot_Utility;
use Limb_Chatbot\Includes\Utilities\Copilot_Utility;

/**
 * Handles LLM‑based reranking of knowledge base entries.
 *
 * This service:
 * - Builds a compact, JSON-encoded payload with the latest user question,
 *   short conversation context and candidate snippets.
 * - Asks the LLM to return ONLY the IDs of relevant candidates as a JSON array.
 * - Validates and normalizes the LLM output, with safe fallbacks.
 */
class Knowledge_Reranker_Service {

	private Chatbot_Utility $chatbot_utility;
	private ?Copilot_Utility $copilot_utility = null;

	public function __construct( Chatbot_Utility $chatbot_utility ) {
		$this->chatbot_utility = $chatbot_utility;
	}

	/**
	 * Reranks the given dataset entries using an LLM.
	 *
	 * @param  Dataset_Entry[]  $entries  Candidate entries from the vector searcher.
	 * @param  Message  $rewritten_message  Past rewritten message (for context).
	 *
	 * @return Dataset_Entry[] Reranked entries. On failure, falls back to the original list.
	 */
	public function rerank( array $entries, Message $rewritten_message ): array {
		if ( count( $entries ) <= 1 ) {
			return $entries;
		}

		$latest_user_message = $rewritten_message->extract_text();
		if ( $latest_user_message === '' ) {
			return $entries;
		}

		// Build compact candidate list with stable numeric IDs and short snippets.
		$candidates = [];
		foreach ( $entries as $entry ) {
			if ( ! $entry instanceof Dataset_Entry ) {
				continue;
			}

			$id = $entry->get_id();
			if ( ! $id ) {
				continue;
			}
			$string = $entry->stringify_for_inference();
			$candidates[] = [
				'id'   => $id,
				'text' => strlen($string) > 1500 ? substr($string, 0, 1500) . '...' : $string,
			];
		}

		if ( empty( $candidates ) ) {
			return $entries;
		}

		// Structured payload keeps things readable in logs and clearer for the LLM.
		$user_payload = json_encode(
			[
				'question'   => $latest_user_message,
				'candidates' => $candidates,
			],
			JSON_UNESCAPED_UNICODE | JSON_UNESCAPED_SLASHES
		);

		try {
			$response = $this->get_copilot_utility()->rerank( $user_payload );

			if ( ! $response instanceof Message ) {
				return $entries;
			}

			$raw = trim( $response->extract_text() );
			if ( $raw === '' ) {
				return $entries;
			}

			$candidate_ids = array_column( $candidates, 'id' );
			$ordered_ids   = $this->parse_reranked_ids( $raw, $candidate_ids );

			if ( empty( $ordered_ids ) ) {
				// Reranker response could not be aligned with our candidates – keep original ranking.
				return $entries;
			}

			$id_map = [];
			foreach ( $entries as $entry ) {
				if ( $entry instanceof Dataset_Entry ) {
					$id_map[ $entry->get_id() ] = $entry;
				}
			}

			$reranked = [];
			foreach ( $ordered_ids as $id ) {
				if ( isset( $id_map[ $id ] ) ) {
					$reranked[] = $id_map[ $id ];
				}
			}

			if ( empty( $reranked ) ) {
				// We received some IDs but none matched existing candidates – safest is to keep original order.
				return $entries;
			}

			return $reranked;
		} catch ( \Throwable $e ) {
			return $entries;
		}
	}

	/**
	 * Lazily builds and returns the Copilot utility used for LLM calls.
	 *
	 * @return Copilot_Utility
	 */
	private function get_copilot_utility(): Copilot_Utility {
		if ( $this->copilot_utility instanceof Copilot_Utility ) {
			return $this->copilot_utility;
		}

		$ai_model       = $this->chatbot_utility->get_ai_model();
		$config         = $this->chatbot_utility->get_config();
		$ai_provider_id = $this->chatbot_utility->get_ai_provider_id();

		if ( ! $ai_model ) {
			throw new \RuntimeException( 'AI model is not configured for knowledge reranking.' );
		}

		$this->copilot_utility = new Copilot_Utility(
			$ai_model,
			$config,
			$ai_provider_id
		);

		return $this->copilot_utility;
	}

	/**
	 * Parses the LLM reranker output into a clean, ordered list of candidate IDs.
	 *
	 * Robustness features:
	 * - Accepts plain JSON arrays like [1734, 1773].
	 * - If the model wraps JSON in extra text, extracts the last numeric-only [ ... ] block.
	 * - As a last resort, collects all numbers from the text and filters by known candidate IDs.
	 *
	 * @param  string  $raw_response  Raw text returned by the LLM.
	 * @param  int[]  $candidate_ids  Allowed candidate IDs.
	 *
	 * @return int[] Clean, de-duplicated list of candidate IDs in model‑preferred order.
	 */
	private function parse_reranked_ids( string $raw_response, array $candidate_ids ): array {
		$raw_response = trim( $raw_response );

		if ( $raw_response === '' ) {
			return [];
		}

		// If the model wrapped the output in a markdown code block (```json ... ```),
		// extract the content inside the fences before further processing.
		if ( preg_match( '/```(?:json)?\s*(.*?)```/is', $raw_response, $code_matches ) ) {
			$raw_response = trim( $code_matches[1] );
		}

		// Helper to normalize an arbitrary decoded JSON value into an array of ints.
		$normalize_ids = static function ( $decoded ) {
			if ( ! is_array( $decoded ) ) {
				return [];
			}

			return array_values(
				array_filter(
					array_map( 'intval', $decoded ),
					static fn( $id ) => $id > 0
				)
			);
		};

		// 1) Happy path: response is exactly a JSON array.
		$ordered_ids = $normalize_ids( json_decode( $raw_response, true ) );

		// 2) If that fails, try to extract the last numeric-only JSON array substring,
		//    e.g. [1770, 1771, 1734]. This avoids picking candidate arrays like [{"id":1770,...}].
		if ( empty( $ordered_ids ) ) {
			if ( preg_match_all( '/\[\s*\d+(?:\s*,\s*\d+)*\s*\]/', $raw_response, $matches ) ) {
				$numeric_arrays = $matches[0] ?? [];
				if ( ! empty( $numeric_arrays ) ) {
					// Use the last numeric array; models often append the final answer at the end.
					$json_fragment = end( $numeric_arrays );
					$ordered_ids   = $normalize_ids( json_decode( $json_fragment, true ) );
				}
			}
		}

		// 3) Final fallback: pull all integers from the text.
		if ( empty( $ordered_ids ) ) {
			if ( preg_match_all( '/\d+/', $raw_response, $matches ) ) {
				$ordered_ids = $normalize_ids( $matches[0] ?? [] );
			}
		}

		if ( empty( $ordered_ids ) || empty( $candidate_ids ) ) {
			return [];
		}

		// Keep only IDs that correspond to our current candidate set.
		$allowed_map = array_fill_keys( array_map( 'intval', $candidate_ids ), true );

		$filtered = [];
		foreach ( $ordered_ids as $id ) {
			$id = (int) $id;
			if ( $id > 0 && isset( $allowed_map[ $id ] ) ) {
				$filtered[] = $id;
			}
		}

		// Preserve order, drop duplicates.
		return array_values( array_unique( $filtered ) );
	}
}
