<?php

namespace Limb_Chatbot\Includes\Services\Knowledge;

use Limb_Chatbot\Includes\Data_Objects\Dataset_Entry;
use Limb_Chatbot\Includes\Traits\SingletonTrait;
use Limb_Chatbot\Includes\Utilities\Chatbot_Utility;

/**
 * Class Knowledge_Source_Store
 *
 * Singleton store for managing knowledge sources during the chatbot conversation flow.
 * Acts as temporary memory that collects sources from various points and provides them
 * when the final AI response is saved.
 *
 * @package Limb_Chatbot\Includes\Services
 * @since 1.0.0
 */
class Knowledge_Source_Store {

	use SingletonTrait;

	/**
	 * Current knowledge sources stored in memory.
	 *
	 * @var Dataset_Entry[]
	 * @since 1.0.0
	 */
	private array $sources = [];

	/**
	 * Whether the sources were set by function calling (which should replace existing).
	 *
	 * @var bool
	 * @since 1.0.0
	 */
	private bool $set_by_function_calling = false;

	/**
	 * Optional chatbot utility for context-aware operations (e.g., distillation).
	 *
	 * @var Chatbot_Utility|null
	 */
	private ?Chatbot_Utility $chatbot_utility = null;

	/**
	 * Latest user question text associated with current knowledge.
	 *
	 * @var string
	 */
	private string $user_question = '';

	/**
	 * Add knowledge sources from User_Prompt_Manager (initial sources).
	 *
	 * These sources are added if no function calling has occurred yet.
	 *
	 * @param  Dataset_Entry[]  $sources  Array of dataset entries.
	 *
	 * @return void
	 * @since 1.0.0
	 */
	public function add_initial_sources( array $sources ): void {
		// Only add if not already overridden by function calling
		if ( ! $this->set_by_function_calling ) {
			$this->sources = $sources;
		}
	}

	/**
	 * Set knowledge sources from function calling (replaces existing sources).
	 *
	 * When function calling occurs, it should replace all existing sources
	 * with the new ones returned by knowledge search.
	 *
	 * @param  Dataset_Entry[]  $sources  Array of dataset entries from function calling.
	 *
	 * @return void
	 * @since 1.0.0
	 */
	public function set_function_calling_sources( array $sources ): void {
		$this->sources                 = $sources;
		$this->set_by_function_calling = true;
	}

	/**
	 * Attach chatbot context and user question for downstream knowledge formatting.
	 *
	 * @param  Chatbot_Utility  $chatbot_utility
	 * @param  string           $user_question
	 *
	 * @return void
	 */
	public function set_context( Chatbot_Utility $chatbot_utility, string $user_question ): void {
		$this->chatbot_utility = $chatbot_utility;
		$this->user_question   = $user_question;
	}

	/**
	 * Get current knowledge sources.
	 *
	 * @return Dataset_Entry[] Array of dataset entries.
	 * @since 1.0.0
	 */
	public function get_sources(): array {
		return $this->sources;
	}

	/**
	 * Check if there are any sources stored.
	 *
	 * @return bool True if sources exist, false otherwise.
	 * @since 1.0.0
	 */
	public function has_sources(): bool {
		return ! empty( $this->sources );
	}

	/**
	 * Get the number of stored sources.
	 *
	 * @return int Number of sources.
	 * @since 1.0.0
	 */
	public function count_sources(): int {
		return count( $this->sources );
	}

	/**
	 * Check if sources were set by function calling.
	 *
	 * @return bool True if set by function calling, false if initial sources.
	 * @since 1.0.0
	 */
	public function is_set_by_function_calling(): bool {
		return $this->set_by_function_calling;
	}

	/**
	 * Consume and clear all stored sources.
	 *
	 * This method is called when the final AI response is saved to extract
	 * the sources and clear the store for the next conversation turn.
	 *
	 * @return Dataset_Entry[] Array of dataset entries that were stored.
	 * @since 1.0.0
	 */
	public function consume_sources(): array {
		$sources = $this->sources;
		$this->clear();

		return $sources;
	}

	/**
	 * Clear all stored sources and reset state.
	 *
	 * @return void
	 * @since 1.0.0
	 */
	public function clear(): void {
		$this->sources                 = [];
		$this->set_by_function_calling = false;
		$this->chatbot_utility        = null;
		$this->user_question          = '';
	}

	/**
	 * Get debug information about the current state.
	 *
	 * @return array Debug information.
	 * @since 1.0.0
	 */
	public function get_debug_info(): array {
		return [
			'source_count'            => count( $this->sources ),
			'set_by_function_calling' => $this->set_by_function_calling,
			'source_ids'              => array_map( function ( $source ) {
				return $source->get_id();
			}, $this->sources ),
		];
	}

	public function stringify_knowledge( ?array $knowledge = null, int $max_chunks = 10 ): string {
		$knowledge = $knowledge ?? $this->sources ?? [];
		if ( empty( $knowledge ) ) {
			return "No references available.";
		}

		$chunks = [];
		$qas    = [];

		foreach ( $knowledge as $entry ) {
			$metadata = $entry->get_entry()['entry_metadata'] ?? [];
			$type     = $metadata['entry_type'] ?? 'chunk';
			if ( $type === 'chunk' ) {
				$chunks[] = $entry;
			} else {
				$qas[] = $entry;
			}
		}

		// Respect max_chunks (especially important in short mode)
		$chunks = array_slice( $chunks, 0, $max_chunks );

		$formatted = [];

		// Chunks: Clean, numbered references
		foreach ( $chunks as $i => $entry ) {
			$idx         = $i + 1;
			$formatted[] = "[Reference {$idx}]\n" . trim( $entry->stringify_for_inference() );
		}

		// QA entries: Clear Q&A format
		foreach ( $qas as $i => $entry ) {
			$idx         = count( $chunks ) + $i + 1;
			$formatted[] = "[QA {$idx}]\n" . trim( $entry->stringify_for_inference() );
		}

		return implode( "\n\n", $formatted );
	}
}
