<?php
/**
 * Context Base Class.
 *
 * Centralizes the management of conversation context for LLM interactions.
 * Handles chronological message log including tool calls and results. Stores
 * tool calls and results chronologically in chat log to maintain proper
 * conversation flow for LLMs. Extended by provider-specific context
 * implementations.
 *
 * @package   Agenistrator
 * @since     1.0.0
 * @author    Alex Ruban <alex@maximum.guru>
 * @author    Maximum Software Inc <contact@maximum.software>
 * @license   GPL-3.0-only
 * @link      https://maximum.software/
 *
 * SPDX-FileCopyrightText: 2025-2026 Alex Ruban <alex@maximum.guru>
 * SPDX-FileCopyrightText: 2025-2026 Maximum Software Inc <contact@maximum.software>
 * SPDX-License-Identifier: GPL-3.0-only
 */

namespace Agenistrator;

// Direct access prevention
if ( ! defined( 'ABSPATH' ) ) {
	return;
}

/**
 * Context class.
 */
abstract class Context {
	/**
	 * System message at the beginning of each context
	 * @var System_Message
	 * @since     1.0.0
	 */
	protected $system_message;

	/**
	 * Messages array containing Message instances
	 * All messages are stored here chronologically
	 * @var Message[]
	 * @since     1.0.0
	 */
	protected $messages = array();

	/**
	 * Context-wide relevant information (injected after system message)
	 * @var string
	 * @since     1.0.0
	 */
	protected $relevant_information = '';

	/**
	 * Available tools for this context (Tool instances)
	 * @var Tool[]
	 * @since     1.0.0
	 */
	protected $available_tools = array();

	/**
	 * Token counter instance (lazy loaded)
	 * @var Token_Counter|null
	 */
	private $token_counter = null;

	/**
	 * Cached token counts keyed by type (system_message, relevant_info, tools)
	 * @var array<string, array{count: int, encoding: string}>
	 */
	private $token_cache = array();

	/**
	 * Get the token counter instance
	 *
	 * @return Token_Counter
	 * @since     1.0.0
	 */
	protected function token_counter() {
		if ( null === $this->token_counter ) {
			$this->token_counter = Token_Counter::get_instance();
		}
		return $this->token_counter;
	}

	/**
	 * Constructor
	 *
	 * @param      string     $system_message Optional custom system message
	 * @since     1.0.0
	 */
	public function __construct( $system_message = null ) {
		$content              = $system_message ?? $this->get_system_prompt();
		$this->system_message = new System_Message( $content );
	}

	/**
	 * Get the system prompt
	 *
	 * @return string The system prompt
	 */
	private function get_system_prompt() {
		return <<<'PROMPT'
You are a helpful agentic AI assistant called 'Agenistrator' that runs inside a WordPress admin dashboard.
You act on the user's behalf, eliminating repetitive work and dramatically accelerating user's workflow.
Instead of hunting through menus and settings pages or getting bogged down with manual data entry, the user will tell you what they need done and you will handles the rest.
You have access to a large array of tools that you can use to complete user tasks, which include tools that are not related to WordPress.
You MUST first call the `list_all_tools` tool and analyzing its output to find out what tools are available before responding to the user that you can't complete their task.

## Language and Locale

Use the user language/locale for all responses, unless the user starts talking to you in a different language or explicitly asks you to use a different language, then use the requested language.

## Message Formatting Requirements

- User messages should be interpreted as markdown-formatted text
- ALL assistant responses (except tool call arguments unless explicitly required by tool) MUST use proper markdown formatting including:
  - **Bold** and *italic* text for emphasis
  - Bulleted lists using `-` or `*` and numbered lists using `1.`
  - Links using proper markdown syntax: [link text](URL)
  - Images using markdown syntax: ![alt text](URL) - use this to display images to the user when relevant (e.g., media library items, screenshots, generated images). The URL must be an absolute URL accessible from the user's browser (not a file path or internal URL).
  - Tables using markdown table syntax when presenting structured data
  - GitHub Flavored Markdown extensions: ~~strikethrough~~, task lists with `- [ ]` and `- [x]`
  - Code blocks with syntax highlighting: ```language for multi-line code and `backticks` for inline code
  - Headers using `#`, `##`, `###` etc. as appropriate
  - Blockquotes using `>` for important information or quotes
- UTF-8 icons and emojis are encouraged to enhance communication
- Always format responses to be visually appealing and easy to read
- Markdown is a superset of a small subset of HTML. Escape anything that looks like HTML but is not meant to be rendered as HTML. Avoid using HTML tags directly.
- When displaying HTML/XML code to the user, ALWAYS wrap them in backticks (inline) or code blocks (multi-line).

## Tool Usage

- You MUST use the `list_all_tools` tool to find out what tools are available.
- You MUST NOT tell the user that you can't complete a task without first calling the `list_all_tools` tool and analyzing its output.
- You MUST import tools that you want to use using the `import_tools` tool before you can use them.
- You MUST unimport previously imported tools once you are done with all tool calls and BEFORE giving the final response.
- You are encouraged to batch multiple tool calls in a single response.

## List Tools, Pagination, and Content Retrieval

- **List tools return summaries only:** All tools that list items (posts, pages, comments, categories, tags, plugins, themes, menus, roles, shortcodes, etc.) return truncated excerpts and descriptions limited to 200 bytes. This prevents information overload.
- **Use content retrieval tools for full content:** To get complete content for any item, use the corresponding detail/retrieval tool (e.g., `get_post_details`, `get_page_details`, `get_comment_details`).
- **List pagination:** All list tools support pagination using `limit` and `offset` parameters. The response includes:
  - The actual pagination parameters used: `limit` and `offset` (when applicable)
  - `has_more`: Boolean flag indicating whether more items exist beyond the current page
  - To retrieve the next page, call the tool again with `offset` incremented by `limit`
- **Content chunking:** Content retrieval tools return large content in chunks to avoid information overload. Use `content_offset` (default: 0) and `content_max_size` (default: 8192 bytes) parameters to control chunking. Note that `content_max_size` is a suggestion only - the actual chunk size may be larger because chunks are aligned to token boundaries to prevent breaking tokens mid-way. The response includes:
  - `content`: The current chunk
  - `content_offset`: Current byte offset
  - `content_size`: Actual size of current chunk in bytes
  - `content_total_size`: Total content size in bytes
  - `content_has_more`: Boolean indicating if more content exists
- **Navigating chunks:** To retrieve the next chunk, use `content_offset + content_size` from the previous response as the new `content_offset` parameter. Continue until `content_has_more` is false.

## Task Execution Planning

When the user asks you to perform a task, you MUST take the following steps:
1. Call `list_all_tools` to find out what tools are available.
2. Look through the list of tools, strategize, and create a TODO list that contains a multi-step plan for completing the task.
3. Import the tools that are needed to execute the plan using the `import_tools` tool.
4. Continuously execute the plan by calling the tools in the order you decided on.
5. Continuously respond with an updated TODO list as you call tools to track your progress and keep the user informed about the progress of the task. If you have no other tool to call but want to share progress, call the `continue` tool in conjunction with your message response.
If any tool call fails, try to work around it by doing anything up to and including restarting from step 1.

## Handling Special Cases

- **Giving manual instructions:** Don't tell the user how to perform the task manually if it can be done with tools that you have access to.
- You MUST use the `search_content` tool to find relevant site content after each user message, unless no meaningful search query can be composed (e.g., greetings, meta questions about your capabilities). This helps you understand context and provide informed responses.
- You MUST use the `calculate_php` tool in every situation where you can write a PHP script to compute an answer (involving math, character and string manipulation, binary logic, etc.). Don't try to guess the answer. The `calculate_php` tool needs a `code` parameter containing your script.
- If `list_all_tools` returns no tools that can be used to answer a user question or satisfy a user request, you MUST respond with a message saying that you don't know how to answer the question or perform the requested action and that you are waiting for further instructions.
- Before writing code and reinventing the wheel, first try to find an already existing plugin (using the `search_plugins` tool) that can help you achieve the desired outcome, however, you must analyze the plugin's capabilities first to make sure it doesn't need any manual user configuration to be used. Plugin web page can be fetched from `https://wordpress.org/plugins/{plugin_slug}/`.

## Multi-Step Task Execution

You are an autonomous agent. Keep working until the user's request is completely resolved before ending your turn and yielding back to the user. Only end your turn when you are sure the task is finished or you need user input.

HOW THE SYSTEM WORKS:
- When you include tool calls in your response, those tools execute and you continue working.
- When you respond WITHOUT tool calls, your turn ends and control returns to the user.

RULES:
- While working on a task, ALWAYS include at least one tool call to keep your turn going.
- When sharing intermediate progress or reasoning, include the `continue` tool alongside your message.
- When the task is COMPLETE, provide a final summary to the user WITHOUT any tool calls. This naturally ends your turn.
- Alternatively, you can call the `terminate` tool to explicitly end your turn (useful after cleanup like `unimport_tools`).
PROMPT;
	}

	/**
	 * Get instruction reminders for synthetic user message
	 *
	 * These reminders reinforce key behaviors that might be lost over long conversations.
	 * They are injected as a synthetic user message after the last real user message to leverage:
	 * - Instruction hierarchy: user messages have higher priority than tool results
	 * - Recency bias: closer to the model's output than the original user message
	 *
	 * @return string Instruction reminders text
	 * @since     1.0.0
	 */
	protected function get_instruction_reminders() {
		return <<<'REMINDERS'
## Reminders
- PERSISTENCE: Keep going until the task is FULLY resolved. Include tool calls to continue your turn.
- PLANNING: Reflect on previous actions before your next step. Maintain and update your TODO list.
- VERIFY: Use tools (`search_content`, `calculate_php`) to verify - NEVER guess or fabricate.
- TOOLS: Call `list_all_tools` to discover tools. Import before use, unimport when done.
REMINDERS;
	}

	/**
	 * Get the header text for relevant information injection
	 *
	 * @return string The header text
	 */
	private function get_relevant_info_header() {
		return __( 'The following information may be relevant for your response:', 'agenistrator' );
	}

	/**
	 * Set the system message
	 *
	 * @param      string     $message System message content
	 * @return     Context    This instance for method chaining
	 * @since     1.0.0
	 */
	public function set_system_message( $message ) {
		$this->system_message = new System_Message( $message );
		unset( $this->token_cache['system_message'] );
		return $this;
	}

	/**
	 * Get the system message
	 *
	 * @return     System_Message     The system message instance
	 * @since     1.0.0
	 */
	public function get_system_message() {
		return $this->system_message;
	}

	/**
	 * Load persistent messages from database into the messages array
	 *
	 * @param Thread $thread The thread instance.
	 * @return Context This instance for method chaining
	 * @since     1.0.0
	 */
	public function load_persistent_messages( Thread $thread ) {
		$messages       = Chat_Log_Manager::get_instance()->get_chat_log( $thread->get_id() );
		$this->messages = array_merge( $this->messages, $messages );
		return $this;
	}

	/**
	 * Add a message to the context
	 *
	 * @param Message $message The message instance to add
	 * @return Context This instance for method chaining
	 * @since     1.0.0
	 */
	public function add_message( Message $message ) {
		$this->messages[] = $message;
		return $this;
	}

	/**
	 * Get all messages
	 *
	 * @return Message[] Array of message instances
	 * @since     1.0.0
	 */
	public function get_messages() {
		return $this->messages;
	}

	/**
	 * Set messages array (used by Context_Pruner for pruned contexts)
	 *
	 * @param Message[] $messages Array of Message instances
	 * @return Context This instance for method chaining
	 * @since     1.0.0
	 */
	public function set_messages( array $messages ) {
		$this->messages = $messages;
		return $this;
	}

	/**
	 * Clear all messages
	 *
	 * @return Context This instance for method chaining
	 * @since     1.0.0
	 */
	public function clear_messages() {
		$this->messages = array();
		return $this;
	}

	/**
	 * Add a user message
	 *
	 * @param string $content Message content
	 * @return Context This instance for method chaining
	 * @since     1.0.0
	 */
	public function add_user_message( $content ) {
		$this->messages[] = new User_Message( $content );
		return $this;
	}

	/**
	 * Add an assistant message
	 *
	 * @param string $content Message content
	 * @return Context This instance for method chaining
	 * @since     1.0.0
	 */
	public function add_assistant_message( $content ) {
		$this->messages[] = new Assistant_Message( $content );
		return $this;
	}

	/**
	 * Add a system message to the messages array
	 *
	 * @param string $content Message content
	 * @return Context This instance for method chaining
	 * @since     1.0.0
	 */
	public function add_system_message( $content ) {
		$this->messages[] = new System_Message( $content );
		return $this;
	}

	/**
	 * Add a tool call message
	 *
	 * @param string $tool_id The ID of the tool being called
	 * @param array $arguments The tool arguments as an associative array
	 * @param string $call_id Required call ID from LLM provider
	 * @return Context This instance for method chaining
	 * @since     1.0.0
	 */
	public function add_tool_call_message( $tool_id, $arguments, $call_id ) {
		$this->messages[] = new Tool_Call_Message( $tool_id, $arguments, $call_id );
		return $this;
	}

	/**
	 * Add a tool result message
	 *
	 * @param string $call_id Required call ID from LLM provider (must match tool call ID)
	 * @param string $tool_id The ID of the tool that was executed
	 * @param mixed $result The tool execution result
	 * @param bool $success Whether the tool execution was successful (defaults to true)
	 * @return Context This instance for method chaining
	 * @since     1.0.0
	 */
	public function add_tool_result_message( $call_id, $tool_id, $result, $success = true ) {
		$this->messages[] = new Tool_Result_Message( $call_id, $tool_id, $result, $success );
		return $this;
	}

	/**
	 * Set context-wide relevant information (injected after system message)
	 *
	 * @param string $information Context-wide relevant information
	 * @return Context This instance for method chaining
	 * @since     1.0.0
	 */
	public function set_relevant_information( $information ) {
		if ( $this->relevant_information !== $information ) {
			$this->relevant_information = $information;
			unset( $this->token_cache['relevant_info'] );
		}
		return $this;
	}

	/**
	 * Get the context-wide relevant information
	 *
	 * @return string The relevant information
	 * @since     1.0.0
	 */
	public function get_relevant_information() {
		return $this->relevant_information;
	}

	/**
	 * Get relevant information from message metadata
	 *
	 * @param Message $message The message to check
	 * @return string|null The relevant information or null if not present
	 */
	private function get_message_relevant_info( Message $message ) {
		$metadata = $message->get_metadata();
		return is_array( $metadata ) && ! empty( $metadata['relevant_information'] )
			? $metadata['relevant_information']
			: null;
	}

	/**
	 * Token overhead per message (approximate)
	 */
	protected const MESSAGE_OVERHEAD         = 4;
	protected const FIRST_MESSAGE_OVERHEAD   = 7;
	protected const FINAL_ASSISTANT_OVERHEAD = 3;

	/**
	 * Token overhead for tool definitions
	 */
	protected const TOOL_FUNC_INIT = 7;
	protected const TOOL_PROP_INIT = 3;
	protected const TOOL_PROP_KEY  = 3;
	protected const TOOL_ENUM_INIT = -3;
	protected const TOOL_ENUM_ITEM = 3;
	protected const TOOL_FUNC_END  = 12;

	/**
	 * Token overhead for tool call/result messages
	 */
	protected const TOOL_CALL_OVERHEAD   = 3;
	protected const TOOL_RESULT_OVERHEAD = 3;

	/**
	 * Count total tokens for the entire context
	 *
	 * Uses cached token counts for system message, relevant info, tools, and message metadata.
	 * Per-message relevant info tokens are included in message token counts.
	 * Cache is invalidated when content changes via setters.
	 *
	 * @return array{total: int, message_starts: int[]} Total token count and message start positions
	 * @since     1.0.0
	 */
	public function count_tokens() {
		$position       = 0;
		$message_starts = array();

		// Tool definitions (cached)
		$position += $this->get_tools_tokens();

		// System message (cached)
		if ( null !== $this->system_message ) {
			$position += $this->get_system_message_tokens();
			$position += static::FIRST_MESSAGE_OVERHEAD;
		}

		// Context-wide relevant information as simulated tool call/result (cached, never pruned)
		$relevant_info_tokens = $this->get_relevant_info_tokens();
		if ( $relevant_info_tokens > 0 ) {
			// Tool call message (assistant with tool_calls)
			$position += $this->get_relevant_info_tool_call_tokens();
			$position += static::MESSAGE_OVERHEAD;
			// Tool result message
			$position += $relevant_info_tokens;
			$position += static::TOOL_RESULT_OVERHEAD;
			$position += static::MESSAGE_OVERHEAD;
		}

		// Find last user message index for instruction reminders
		$last_user_index = null;
		foreach ( $this->messages as $i => $message ) {
			if ( $message instanceof User_Message ) {
				$last_user_index = $i;
			}
		}

		// Conversation messages - record start position of each
		// Per-message relevant info tokens and overhead are included in message token counts
		foreach ( $this->messages as $i => $message ) {
			$message_starts[ $i ] = $position;
			$position            += $this->get_message_tokens( $message );
			$position            += static::MESSAGE_OVERHEAD;

			// Synthetic instruction reminders message after last user message
			if ( $i === $last_user_index ) {
				$position += $this->get_instruction_reminders_tokens();
				$position += static::MESSAGE_OVERHEAD;
			}
		}
		$message_starts[] = $position; // Position after last message

		// Final assistant message overhead
		$position += static::FINAL_ASSISTANT_OVERHEAD;

		return array(
			'total'          => $position,
			'message_starts' => $message_starts,
		);
	}

	/**
	 * Get token count for system message, using cache when available
	 */
	private function get_system_message_tokens() {
		return $this->get_cached_content_tokens(
			'system_message',
			$this->system_message->get_content()
		);
	}

	/**
	 * Get token count for context-wide relevant information, using cache when available
	 */
	private function get_relevant_info_tokens() {
		if ( empty( $this->relevant_information ) ) {
			return 0;
		}
		return $this->get_cached_content_tokens(
			'relevant_info',
			$this->get_relevant_info_header() . "\n\n" . $this->relevant_information
		);
	}

	/**
	 * Get token count for the synthetic tool call used to inject relevant information
	 *
	 * Approximates tokens for: tool name, empty arguments "{}", and call_id pattern.
	 * This is a fixed overhead since the tool call structure is constant.
	 *
	 * @return int Token count for the synthetic tool call
	 */
	private function get_relevant_info_tool_call_tokens() {
		// Tool name "_get_relevant_information" + empty args "{}" + call_id "_ctx_0" or "_msg_N"
		// Approximate: ~15 tokens (tool name ~5, args ~2, call_id ~3, structure ~5)
		return 15 + static::TOOL_CALL_OVERHEAD;
	}

	/**
	 * Get token count for instruction reminders message, using cache when available
	 */
	private function get_instruction_reminders_tokens() {
		return $this->get_cached_content_tokens(
			'instruction_reminders',
			$this->get_instruction_reminders()
		);
	}

	/**
	 * Get token count for content with caching
	 */
	private function get_cached_content_tokens( $key, $content ) {
		$token_counter = $this->token_counter();
		$encoding      = $token_counter->get_encoding();
		if ( isset( $this->token_cache[ $key ] ) && $this->token_cache[ $key ]['encoding'] === $encoding ) {
			return $this->token_cache[ $key ]['count'];
		}
		$count                     = $token_counter->count_tokens( $content );
		$this->token_cache[ $key ] = array(
			'count'    => $count,
			'encoding' => $encoding,
		);
		return $count;
	}

	/**
	 * Get token count for tool definitions, using cache when available
	 *
	 * Tool definitions consume tokens from the context window and must be counted.
	 *
	 * Algorithm (for each tool):
	 * - Add TOOL_FUNC_INIT (7) tokens overhead
	 * - Encode "name:description" (strip trailing period from description)
	 * - If has properties: add TOOL_PROP_INIT (3) tokens
	 * - For each property: add TOOL_PROP_KEY (3), encode "name:type:description"
	 * - For enums: add TOOL_ENUM_INIT (-3), then per item add TOOL_ENUM_ITEM (3) + encode value
	 * - After all tools: add TOOL_FUNC_END (12) tokens
	 *
	 * @see https://cookbook.openai.com/examples/how_to_count_tokens_with_tiktoken
	 *
	 * @return int Token count for all tool definitions
	 */
	private function get_tools_tokens() {
		$token_counter = $this->token_counter();
		$encoding      = $token_counter->get_encoding();
		if ( isset( $this->token_cache['tools'] ) && $this->token_cache['tools']['encoding'] === $encoding ) {
			return $this->token_cache['tools']['count'];
		}

		$count = 0;
		if ( empty( $this->available_tools ) ) {
			$this->token_cache['tools'] = array(
				'count'    => 0,
				'encoding' => $encoding,
			);
			return 0;
		}

		foreach ( $this->available_tools as $tool ) {
			$count += self::TOOL_FUNC_INIT;

			// Tool name, description, description_instructions, and instructions (strip trailing period)
			$desc = $tool->get_description();
			if ( String_Utils::mb_substr( $desc, -1 ) === '.' ) {
				$desc = String_Utils::mb_substr( $desc, 0, -1 );
			}
			$description_instructions = $tool->get_description_instructions();
			if ( '' !== $description_instructions ) {
				$desc .= ' ' . $description_instructions;
			}
			$instructions = $tool->get_instructions();
			if ( '' !== $instructions ) {
				$desc .= "\n\n" . $instructions;
			}
			$count += $token_counter->count_tokens( $tool->get_id() . ':' . $desc );

			// Parameters
			$params = $tool->get_parameters();
			if ( ! empty( $params['properties'] ) ) {
				$count += self::TOOL_PROP_INIT;
				foreach ( $params['properties'] as $prop_name => $prop ) {
					$count += self::TOOL_PROP_KEY;

					// Property name, type, and description
					$prop_type = $prop['type'] ?? 'string';
					$prop_desc = $prop['description'] ?? '';
					if ( String_Utils::mb_substr( $prop_desc, -1 ) === '.' ) {
						$prop_desc = String_Utils::mb_substr( $prop_desc, 0, -1 );
					}
					$count += $token_counter->count_tokens( $prop_name . ':' . $prop_type . ':' . $prop_desc );

					// Enum values
					if ( ! empty( $prop['enum'] ) ) {
						$count += self::TOOL_ENUM_INIT;
						foreach ( $prop['enum'] as $enum_value ) {
							$count += self::TOOL_ENUM_ITEM;
							$count += $token_counter->count_tokens( $enum_value );
						}
					}
				}
			}
		}

		// Final overhead after all tools
		$count += self::TOOL_FUNC_END;

		// Cache the result
		$this->token_cache['tools'] = array(
			'count'    => $count,
			'encoding' => $encoding,
		);

		return $count;
	}

	/**
	 * Get token count for a message, using cache when available
	 *
	 * @param Message $message The message to count tokens for
	 * @return int Token count for the message content
	 */
	private function get_message_tokens( Message $message ) {
		$encoding = $this->token_counter()->get_encoding();

		// Check cache first
		$cached = $this->get_cached_token_count( $message, $encoding );
		if ( null !== $cached ) {
			return $cached;
		}

		// Compute and cache the result
		$tokens = $this->compute_message_tokens( $message );
		$this->set_cached_token_count( $message, $tokens, $encoding );

		// Persist to database if message is saved
		if ( $message->get_message_id() !== null ) {
			Chat_Log_Manager::get_instance()->update_message_metadata(
				$message->get_message_id(),
				$message->get_metadata()
			);
		}

		return $tokens;
	}

	/**
	 * Get cached token count from message metadata or database
	 *
	 * @param Message $message The message to check
	 * @param string $encoding The tokenizer encoding to match
	 * @return int|null Token count if cached with matching encoding, null otherwise
	 */
	private function get_cached_token_count( Message $message, $encoding ) {
		// Check in-memory metadata first
		$cached = $this->extract_token_count( $message->get_metadata(), $encoding );
		if ( null !== $cached ) {
			return $cached;
		}

		// If message is persisted, check database for fresh metadata (handles stale instances)
		$message_id = $message->get_message_id();
		if ( null === $message_id ) {
			return null;
		}

		$db_metadata = Chat_Log_Manager::get_instance()->get_message_metadata( $message_id );
		$cached      = $this->extract_token_count( $db_metadata, $encoding );
		if ( null !== $cached ) {
			// Sync in-memory metadata with database
			$message->set_metadata( $db_metadata );
		}
		return $cached;
	}

	/**
	 * Extract token count from metadata array if encoding matches
	 */
	private function extract_token_count( $metadata, $encoding ) {
		if ( ! is_array( $metadata ) || ! isset( $metadata['token_count'] ) ) {
			return null;
		}
		$cache = $metadata['token_count'];
		if ( ! isset( $cache['encoding'] ) || $cache['encoding'] !== $encoding ) {
			return null;
		}
		return $cache['count'] ?? null;
	}

	/**
	 * Set cached token count in message metadata
	 *
	 * @param Message $message The message to update
	 * @param int $count The token count
	 * @param string $encoding The tokenizer encoding used
	 */
	private function set_cached_token_count( Message $message, $count, $encoding ) {
		$metadata = $message->get_metadata();
		if ( ! is_array( $metadata ) ) {
			$metadata = array();
		}
		$metadata['token_count'] = array(
			'count'    => $count,
			'encoding' => $encoding,
		);
		$message->set_metadata( $metadata );
	}

	/**
	 * Compute token count for a message (without caching)
	 *
	 * Includes relevant_information tokens from metadata if present.
	 *
	 * @param Message $message The message to count tokens for
	 * @return int Token count
	 * @since     1.0.0
	 */
	protected function compute_message_tokens( Message $message ) {
		$token_counter = $this->token_counter();
		$tokens        = 0;

		$content = $message->get_content();
		if ( null !== $content ) {
			$tokens += $token_counter->count_tokens( $content );
		}

		if ( $message instanceof Tool_Call_Message ) {
			$tokens += $token_counter->count_tokens( $message->get_tool_id() );
			// phpcs:ignore WordPress.WP.AlternativeFunctions.json_encode_json_encode -- Using json_encode() directly to avoid wp_json_encode_options filter interference; token counting requires consistent encoding.
			$tokens += $token_counter->count_tokens( json_encode( $message->get_arguments() ) );
			$call_id = $message->get_call_id();
			if ( null !== $call_id ) {
				$tokens += $token_counter->count_tokens( $call_id );
			}
			$tokens += static::TOOL_CALL_OVERHEAD;
		} elseif ( $message instanceof Tool_Result_Message ) {
			$tokens += $token_counter->count_tokens( $message->get_tool_id() );
			// phpcs:ignore WordPress.WP.AlternativeFunctions.json_encode_json_encode -- Using json_encode() directly to avoid wp_json_encode_options filter interference; token counting requires consistent encoding.
			$tokens += $token_counter->count_tokens( json_encode( $message->get_result() ) );
			$call_id = $message->get_call_id();
			if ( null !== $call_id ) {
				$tokens += $token_counter->count_tokens( $call_id );
			}
			$tokens += static::TOOL_RESULT_OVERHEAD;
		}

		// Include relevant information tokens as simulated tool call/result from metadata
		$relevant_info = $this->get_message_relevant_info( $message );
		if ( null !== $relevant_info ) {
			// Tool call message tokens
			$tokens += $this->get_relevant_info_tool_call_tokens();
			$tokens += static::MESSAGE_OVERHEAD;
			// Tool result message tokens
			$tokens += $token_counter->count_tokens(
				$this->get_relevant_info_header() . "\n\n" . $relevant_info
			);
			$tokens += static::TOOL_RESULT_OVERHEAD;
			$tokens += static::MESSAGE_OVERHEAD;
		}

		return $tokens;
	}

	/**
	 * Set available tools
	 *
	 * @param      array      $tools Array of Tool instances
	 * @return     Context    This instance for method chaining
	 * @throws     \Exception If any tool is not a Tool instance
	 * @since     1.0.0
	 */
	public function set_available_tools( $tools ) {
		// Validate that all tools are instances of Tool
		foreach ( $tools as $tool ) {
			if ( ! ( $tool instanceof Tool ) ) {
				throw new \Exception( 'Invalid tool provided: tools must be instances of Tool' );
			}
		}

		$this->available_tools = $tools;
		unset( $this->token_cache['tools'] );
		return $this;
	}

	/**
	 * Get the available tools
	 *
	 * @return     Tool[]     Array of tool instances
	 * @since     1.0.0
	 */
	public function get_available_tools() {
		return $this->available_tools;
	}

	/**
	 * Render the messages array for the LLM
	 *
	 * Context-wide relevant information is injected after system message.
	 * Per-message relevant information is injected after messages that have it in metadata.
	 *
	 * @return     array      Properly formatted messages array
	 */
	/** Synthetic tool name for injecting relevant information */
	protected const RELEVANT_INFO_TOOL_NAME = '_get_relevant_information';

	/**
	 * Render the messages array for the LLM
	 *
	 * Context-wide relevant information is injected after system message as a simulated tool call/result.
	 * Per-message relevant information is injected after messages that have it, also as tool call/result.
	 * This approach avoids multiple system messages which can dilute core instructions.
	 * Instruction reminders are added as a synthetic user message after the last user message.
	 *
	 * @return array Properly formatted messages array
	 * @since     1.0.0
	 */
	public function render_messages() {
		$formatted_messages = array();

		// Find the last user message index for instruction reminders
		$last_user_index = null;
		foreach ( $this->messages as $i => $message ) {
			if ( $message instanceof User_Message ) {
				$last_user_index = $i;
			}
		}

		// 1. Always start with the system message
		if ( null !== $this->system_message ) {
			$formatted_messages[] = $this->format_message( $this->system_message );
		}

		// 2. Context-wide relevant information (as simulated tool call/result)
		if ( ! empty( $this->relevant_information ) ) {
			$content              = $this->get_relevant_info_header() . "\n\n" . $this->relevant_information;
			$call_id              = '_ctx_0';
			$formatted_messages[] = $this->format_message(
				new Tool_Call_Message( static::RELEVANT_INFO_TOOL_NAME, array(), $call_id )
			);
			$formatted_messages[] = $this->format_message(
				new Tool_Result_Message( $call_id, static::RELEVANT_INFO_TOOL_NAME, $content, true )
			);
		}

		// 3. Process all messages chronologically, injecting per-message relevant info
		// Note: Tool discovery filtering is now handled by Context_Pruner
		foreach ( $this->messages as $i => $message ) {
			$formatted_msg = $this->format_message( $message );
			if ( null !== $formatted_msg ) {
				$formatted_messages[] = $formatted_msg;
			}

			// Inject per-message relevant information after messages that have it (as simulated tool call/result)
			$relevant_info = $this->get_message_relevant_info( $message );
			if ( null !== $relevant_info ) {
				$content              = $this->get_relevant_info_header() . "\n\n" . $relevant_info;
				$call_id              = '_msg_' . $i;
				$formatted_messages[] = $this->format_message(
					new Tool_Call_Message( static::RELEVANT_INFO_TOOL_NAME, array(), $call_id )
				);
				$formatted_messages[] = $this->format_message(
					new Tool_Result_Message( $call_id, static::RELEVANT_INFO_TOOL_NAME, $content, true )
				);
			}

			// Inject instruction reminders as synthetic user message after last user message
			if ( $i === $last_user_index ) {
				$formatted_messages[] = $this->format_message(
					new User_Message( $this->get_instruction_reminders() )
				);
			}
		}

		return $formatted_messages;
	}

	/**
	 * Format a Message instance according to its type
	 * Provider-specific classes should override this method
	 *
	 * @param Message $message The message instance to format
	 * @return array|null Formatted message or null if invalid
	 * @since     1.0.0
	 */
	abstract protected function format_message( Message $message );

	/**
	 * Count tokens for a completion response
	 *
	 * Used for estimating output tokens when API doesn't return usage data.
	 * Provider-agnostic: accepts Message instances so any provider can use it.
	 * Models a single assistant response that may contain content and/or tool calls.
	 *
	 * @param Assistant_Message|null $content Optional assistant message with text content
	 * @param Tool_Call_Message[] $tool_calls Array of tool call messages
	 * @return int Token count for the completion
	 * @since     1.0.0
	 */
	public function count_completion_tokens( $content, $tool_calls = array() ) {
		$tokens = 0;

		// Content tokens
		if ( null !== $content ) {
			$tokens += $this->compute_message_tokens( $content );
		}

		// Tool call tokens (each includes TOOL_CALL_OVERHEAD via compute_message_tokens)
		foreach ( $tool_calls as $tool_call ) {
			$tokens += $this->compute_message_tokens( $tool_call );
		}

		// Single MESSAGE_OVERHEAD for the entire completion response
		$tokens += static::MESSAGE_OVERHEAD;

		return $tokens;
	}
}
