<?php
/**
 * Context Pruner.
 *
 * Prunes conversation context to fit within model's context window.
 * Applies pruning strategies sequentially until context fits. Currently
 * implements hard truncation only; other strategies to be added later.
 *
 * @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://agenistrator.com/
 *
 * 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' ) ) {
	exit;
}

/**
 * Context_Pruner class.
 */
class Context_Pruner {

	use Singleton;

	/**
	 * Prune context to fit within model's context window
	 *
	 * @param Context $context The context to prune
	 * @param LLM_Chat $chat_model The chat model instance
	 * @return Context Pruned context (cloned if modified, original if no pruning needed)
	 * @since     1.0.0
	 */
	public function prune( Context $context, LLM_Chat $chat_model ) {
		$max_tokens = Model_Info::get_instance()->get_available_input( $chat_model );

		// Clone context to work on
		$pruned_context = clone $context;
		$messages       = $pruned_context->get_messages();

		// Always apply tool discovery filtering first (moved from Context class)
		$messages = $this->filter_tool_discovery_messages( $messages );
		$pruned_context->set_messages( $messages );

		$token_data = $pruned_context->count_tokens();

		// Apply hard truncation if context exceeds limit
		if ( $token_data['total'] > $max_tokens ) {
			$pruned_context = $this->hard_truncate( $pruned_context, $max_tokens, $token_data );
		}

		return $pruned_context;
	}

	/**
	 * Filter out list_all_tools call/result pairs that precede an import_tools call
	 *
	 * Once tools are imported, the verbose list_all_tools output is no longer needed.
	 * The LLM can always call list_all_tools again if needed.
	 *
	 * @param Message[] $messages Array of Message objects
	 * @return Message[] Filtered messages
	 */
	private function filter_tool_discovery_messages( array $messages ) {
		$found_import = false;
		for ( $i = count( $messages ) - 1; $i >= 0; $i-- ) {
			$message = $messages[ $i ];
			if ( ! ( $message instanceof Tool_Call_Message || $message instanceof Tool_Result_Message ) ) {
				continue;
			}
			$tool_id = $message->get_tool_id();
			if ( 'import_tools' === $tool_id ) {
				$found_import = true;
			} elseif ( $found_import && 'list_all_tools' === $tool_id ) {
				unset( $messages[ $i ] );
			}
		}
		return array_values( $messages );
	}

	/**
	 * Hard truncate context to fit within token limit
	 *
	 * Uses message start positions to find the cutoff point.
	 * Accounts for fixed overhead (system message, relevant info).
	 * Ensures tool call/result pairs are removed together.
	 *
	 * @param Context $context Context to truncate
	 * @param int $max_tokens Maximum tokens allowed
	 * @param array $token_data Token data from count_tokens()
	 * @return Context Truncated context
	 */
	private function hard_truncate( Context $context, int $max_tokens, array $token_data ) {
		$messages      = $context->get_messages();
		$message_count = count( $messages );
		if ( 0 === $message_count ) {
			return $context;
		}

		$message_starts         = $token_data['message_starts'];
		$fixed_prefix           = $message_starts[0]; // System message tokens
		$fixed_suffix           = $token_data['total'] - $message_starts[ $message_count ]; // Relevant info + final overhead
		$available_for_messages = $max_tokens - $fixed_prefix - $fixed_suffix;

		if ( $available_for_messages <= 0 ) {
			$context->set_messages( array() );
			return $context;
		}

		// Find minimum cutoff where remaining messages fit
		$messages_end = $message_starts[ $message_count ];
		$min_start    = $messages_end - $available_for_messages;
		$cutoff       = $message_count;

		for ( $i = 0; $i < $message_count; $i++ ) {
			if ( $message_starts[ $i ] >= $min_start ) {
				$cutoff = $i;
				break;
			}
		}

		// Extend cutoff to include orphaned tool results
		$cutoff = $this->extend_cutoff_for_tool_pairs( $messages, $cutoff );

		if ( $cutoff > 0 ) {
			$context->set_messages( array_slice( $messages, $cutoff ) );
		}

		return $context;
	}

	/**
	 * Extend cutoff to include any tool results whose calls are being removed
	 *
	 * @param Message[] $messages Messages array
	 * @param int $cutoff Initial cutoff index
	 * @return int Extended cutoff index
	 */
	private function extend_cutoff_for_tool_pairs( array $messages, int $cutoff ) {
		// Collect call IDs from tool calls being removed (before cutoff)
		$removed_call_ids = array();
		$message_count    = count( $messages );
		for ( $i = 0; $i < $cutoff && $i < $message_count; $i++ ) {
			if ( $messages[ $i ] instanceof Tool_Call_Message ) {
				$removed_call_ids[ $messages[ $i ]->get_call_id() ] = true;
			}
		}

		// Extend cutoff to include any tool results matching removed calls
		while ( $cutoff < $message_count ) {
			$msg = $messages[ $cutoff ];
			if ( $msg instanceof Tool_Result_Message && isset( $removed_call_ids[ $msg->get_call_id() ] ) ) {
				++$cutoff;
			} else {
				break;
			}
		}

		return $cutoff;
	}
}