<?php
/**
 * Intelligent Document Chunking Service.
 *
 * @package AISupportAssistant\Documents
 */

namespace AISupportAssistant\Documents;

use AISupportAssistant\Helpers\TextHelper;

if ( ! defined( 'ABSPATH' ) ) {
	exit;
}

/**
 * Class ChunkingService
 */
class ChunkingService {

	/**
	 * Break text into semantic chunks respecting heading, paragraph, and sentence boundaries.
	 *
	 * @param string $text          Raw cleaned document text.
	 * @param int    $chunk_size    Target maximum tokens per chunk.
	 * @param int    $chunk_overlap Overlapping tokens window size.
	 * @return array Array of chunks metadata.
	 */
	public static function chunk( string $text, int $chunk_size = 700, int $chunk_overlap = 100 ): array {
		if ( empty( trim( $text ) ) ) {
			return array();
		}

		// 1. Break text into paragraphs.
		$paragraphs = explode( "\n\n", $text );
		$elements   = array();

		foreach ( $paragraphs as $para ) {
			$para = trim( $para );
			if ( empty( $para ) ) {
				continue;
			}

			// If paragraph starts with heading markdown characters (#) or is very short, treat as heading block.
			if ( 0 === strpos( $para, '#' ) || ( strlen( $para ) < 100 && false === strpos( $para, "\n" ) ) ) {
				$elements[] = array(
					'text'       => $para,
					'is_heading' => true,
				);
			} else {
				// Otherwise split paragraph into sentences.
				$sentences = TextHelper::split_into_sentences( $para );
				foreach ( $sentences as $sentence ) {
					$elements[] = array(
						'text'       => $sentence,
						'is_heading' => false,
					);
				}
			}
		}

		$chunks        = array();
		$element_count = count( $elements );
		$i             = 0;
		$chunk_num     = 1;

		// 2. Sliding window chunking based on approximated token sizes.
		while ( $i < $element_count ) {
			$current_chunk_text = '';
			$current_tokens     = 0;
			$added_indices      = array();

			for ( $k = $i; $k < $element_count; $k++ ) {
				$element_text   = $elements[ $k ]['text'];
				$element_tokens = TextHelper::count_tokens( $element_text );

				// If we have elements in the chunk and adding this one exceeds chunk_size, stop.
				if ( $current_tokens > 0 && ( $current_tokens + $element_tokens ) > $chunk_size ) {
					break;
				}

				$current_chunk_text .= ( '' !== $current_chunk_text ? "\n" : '' ) . $element_text;
				$current_tokens     += $element_tokens;
				$added_indices[]     = $k;
			}

			if ( '' !== trim( $current_chunk_text ) ) {
				$chunks[] = array(
					'chunk_number' => $chunk_num++,
					'chunk_text'   => $current_chunk_text,
					'token_count'  => $current_tokens,
					'page_number'  => null, // Page number estimation can be enriched during extraction if markers exist.
				);
			}

			if ( empty( $added_indices ) ) {
				break;
			}

			// Determine beginning of next chunk.
			$last_added_index = end( $added_indices );
			$next_start       = $last_added_index + 1; // Default to next element.

			// Backtrack to build the overlap window.
			$overlap_tokens = 0;
			for ( $j = $last_added_index; $j >= $i; $j-- ) {
				$overlap_tokens += TextHelper::count_tokens( $elements[ $j ]['text'] );
				if ( $overlap_tokens >= $chunk_overlap ) {
					$next_start = $j;
					break;
				}
			}

			// Prevent infinite loops if backtracking does not advance.
			if ( $next_start <= $i ) {
				$next_start = $last_added_index + 1;
			}

			$i = $next_start;
		}

		return $chunks;
	}
}