<?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 {
		$chunk_size = max( 2, $chunk_size );
		$chunk_overlap = max( 0, min( $chunk_size - 1, $chunk_overlap ) );
		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,
					);
				}
			}
		}

		// Split oversized elements before packing. Count with the same estimator used below.
		$bounded = array();
		foreach ( $elements as $element ) {
			$remaining = $element['text'];
			while ( TextHelper::count_tokens( $remaining ) > $chunk_size ) {
				$low = 1;
				$high = mb_strlen( $remaining, 'UTF-8' );
				while ( $low < $high ) {
					$mid = (int) ceil( ( $low + $high ) / 2 );
					if ( TextHelper::count_tokens( mb_substr( $remaining, 0, $mid, 'UTF-8' ) ) <= $chunk_size ) {
						$low = $mid;
					} else { $high = $mid - 1; }
				}
				$bounded[] = array( 'text' => mb_substr( $remaining, 0, $low, 'UTF-8' ), 'is_heading' => false );
				$remaining = mb_substr( $remaining, $low, null, 'UTF-8' );
			}
			if ( '' !== trim( $remaining ) ) { $bounded[] = array( 'text' => $remaining, 'is_heading' => $element['is_heading'] ); }
		}
		$elements = $bounded;
		$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.
				$candidate = $current_chunk_text . ( '' !== $current_chunk_text ? "\n" : '' ) . $element_text;
				if ( $current_tokens > 0 && TextHelper::count_tokens( $candidate ) > $chunk_size ) {
					break;
				}

				$current_chunk_text .= ( '' !== $current_chunk_text ? "\n" : '' ) . $element_text;
				$current_tokens      = TextHelper::count_tokens( $current_chunk_text );
				$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; $chunk_overlap > 0 && $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;
	}

	/** Keep PDF page boundaries intact so each citation has a real source page. */
	public static function chunk_pages( array $pages, int $size = 700, int $overlap = 100 ): array {
		$result = array();
		foreach ( $pages as $number => $text ) {
			foreach ( self::chunk( $text, $size, $overlap ) as $chunk ) {
				$chunk['chunk_number'] = count( $result ) + 1;
				$chunk['page_number'] = (int) $number;
				$result[] = $chunk;
			}
		}
		return $result;
	}
}