<?php
/**
 * Chunk Extractor.
 *
 * Handles extraction of text chunks from WordPress posts for embedding generation.
 * Implements section-based chunking with markdown conversion and token-aware splitting.
 *
 * @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;

use League\HTMLToMarkdown\HtmlConverter;
use League\HTMLToMarkdown\Converter\TableConverter;

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

/**
 * Chunk_Extractor class.
 */
class Chunk_Extractor {
	use Singleton;

	/**
	 * Default maximum tokens per chunk for embedding models
	 */
	public const DEFAULT_MAX_CHUNK_TOKENS = 8000;

	/**
	 * Token buffer ratio to account for tokenizer differences (10%)
	 */
	private const TOKEN_BUFFER_RATIO = 0.9;

	/**
	 * Minimum tokens per chunk - chunks below this threshold are merged
	 */
	private const MIN_CHUNK_TOKENS = 50;

	/**
	 * Extract chunks from post content for embedding generation.
	 *
	 * Implements section-based chunking with markdown conversion. Content is first
	 * converted to markdown, then split by headings to keep sections together.
	 * Large sections are further split to respect token limits. Small chunks are
	 * merged with adjacent chunks. Includes post metadata (title, type) as first chunk.
	 *
	 * Returns a result array containing:
	 * - 'chunks': Array of chunk objects with 'content', 'source', 'position', 'length'
	 * - 'content_md5': MD5 hash of the rendered markdown content (for change detection)
	 *
	 * Chunk sources:
	 * - 'title': Generated from post title and type
	 * - 'excerpt': From post excerpt
	 * - 'content': From markdown-rendered post content (has position/length)
	 *
	 * @param      \WP_Post $post The post to extract chunks from
	 * @param      int $max_tokens Maximum tokens per chunk (default: 8000)
	 * @return     array Result array with 'chunks' and 'content_md5' keys
	 * @since     1.0.0
	 */
	public function extract_chunks( $post, $max_tokens = self::DEFAULT_MAX_CHUNK_TOKENS ) {
		$chunks        = array();
		$effective_max = (int) ( $max_tokens * self::TOKEN_BUFFER_RATIO );
		$token_counter = Token_Counter::get_instance();

		// Get basic post metadata for context
		$post_type = $post->post_type;

		// Add title as first chunk with minimal context (if not empty)
		if ( ! empty( trim( $post->post_title ) ) ) {
			$title_chunk = String_Utils::replace_tags(
				__( '{type}: {title}', 'agenistrator' ),
				array(
					'type'  => esc_html( ucfirst( $post_type ) ),
					'title' => esc_html( trim( $post->post_title ) ),
				)
			);
			$chunks[]    = array(
				'content'  => $title_chunk,
				'source'   => 'title',
				'position' => null,
				'length'   => null,
			);
		}

		// Add excerpt as chunk (if exists)
		if ( ! empty( trim( $post->post_excerpt ) ) ) {
			$chunks[] = array(
				'content'  => trim( $post->post_excerpt ),
				'source'   => 'excerpt',
				'position' => null,
				'length'   => null,
			);
		}

		// Convert content to markdown
		$markdown    = $this->convert_post_to_markdown( $post );
		$content_md5 = md5( $markdown );

		if ( ! empty( $markdown ) ) {
			// Extract content chunks with position tracking
			$content_chunks = $this->extract_content_chunks( $markdown, $effective_max, $token_counter );
			$chunks         = array_merge( $chunks, $content_chunks );
		}

		return array(
			'chunks'      => $chunks,
			'content_md5' => $content_md5,
		);
	}

	/**
	 * Extract chunks from markdown content with position tracking.
	 *
	 * Uses section-based chunking: splits on markdown headings to keep headings
	 * with their content. Small chunks are merged with adjacent chunks.
	 *
	 * @param      string $markdown The full markdown content
	 * @param      int $effective_max Maximum tokens per chunk (after buffer applied)
	 * @param      Token_Counter $token_counter Token counter instance
	 * @return     array Array of chunk objects with 'content', 'source', 'position', 'length'
	 */
	private function extract_content_chunks( $markdown, $effective_max, $token_counter ) {
		// Split into sections based on markdown headings
		$sections = $this->split_into_sections( $markdown );
		$chunks   = array();

		foreach ( $sections as $section ) {
			$content  = rtrim( $section['content'] );
			$position = $section['position'];

			if ( empty( trim( $content ) ) ) {
				continue;
			}

			// Check if section fits within token limit
			$token_count = $token_counter->count_tokens( $content );
			if ( $token_count <= $effective_max ) {
				$chunks[] = array(
					'content'  => $content,
					'source'   => 'content',
					'position' => $position,
					'length'   => strlen( $content ),
				);
			} else {
				// Split large section into smaller chunks with position tracking
				$sub_chunks = $this->split_text_by_tokens_with_positions(
					$content,
					$position,
					$effective_max,
					$token_counter
				);
				$chunks     = array_merge( $chunks, $sub_chunks );
			}
		}

		// Merge small chunks with adjacent chunks
		return $this->merge_small_chunks( $chunks, $token_counter, $effective_max );
	}

	/**
	 * Split markdown content into sections based on headings.
	 *
	 * Each section starts with a heading (or is content before the first heading)
	 * and includes all content until the next heading of equal or higher level.
	 *
	 * @param      string $markdown The markdown content
	 * @return     array Array of sections with 'content' and 'position' keys
	 */
	private function split_into_sections( $markdown ) {
		$sections = array();

		// Find all heading positions with their levels
		// Matches lines starting with 1-6 # characters followed by space
		preg_match_all( '/^(#{1,6}) /m', $markdown, $matches, PREG_OFFSET_CAPTURE );

		if ( empty( $matches[0] ) ) {
			// No headings - return entire content as single section
			$content = trim( $markdown );
			if ( ! empty( $content ) ) {
				$sections[] = array(
					'content'  => $content,
					'position' => 0,
				);
			}
			return $sections;
		}

		// Extract heading positions
		$heading_positions = array();
		foreach ( $matches[0] as $index => $match ) {
			$heading_positions[] = array(
				'offset' => $match[1],
				'level'  => strlen( $matches[1][ $index ][0] ),
			);
		}

		// Content before first heading (if any)
		$first_heading_pos = $heading_positions[0]['offset'];
		if ( $first_heading_pos > 0 ) {
			$content = trim( substr( $markdown, 0, $first_heading_pos ) );
			if ( ! empty( $content ) ) {
				$sections[] = array(
					'content'  => $content,
					'position' => 0,
				);
			}
		}

		// Extract each section (heading + content until next heading)
		$total_headings = count( $heading_positions );
		for ( $i = 0; $i < $total_headings; $i++ ) {
			$start = $heading_positions[ $i ]['offset'];
			$end   = ( $i + 1 < $total_headings ) ? $heading_positions[ $i + 1 ]['offset'] : strlen( $markdown );

			$content = trim( substr( $markdown, $start, $end - $start ) );
			if ( ! empty( $content ) ) {
				$sections[] = array(
					'content'  => $content,
					'position' => $start,
				);
			}
		}

		return $sections;
	}

	/**
	 * Merge chunks that are below the minimum token threshold.
	 *
	 * Small chunks are merged with the following chunk to avoid fragmentation.
	 * This keeps related content together (e.g., a short heading with its content).
	 *
	 * @param      array $chunks Array of chunk objects
	 * @param      Token_Counter $token_counter Token counter instance
	 * @param      int $max_tokens Maximum tokens per chunk
	 * @return     array Merged array of chunk objects
	 */
	private function merge_small_chunks( $chunks, $token_counter, $max_tokens ) {
		if ( count( $chunks ) <= 1 ) {
			return $chunks;
		}

		$merged  = array();
		$pending = null;

		foreach ( $chunks as $chunk ) {
			if ( null === $pending ) {
				$pending = $chunk;
				continue;
			}

			$pending_tokens = $token_counter->count_tokens( $pending['content'] );

			// If pending chunk is small, try to merge with current chunk
			if ( $pending_tokens < self::MIN_CHUNK_TOKENS ) {
				$combined_content = $pending['content'] . "\n\n" . $chunk['content'];
				$combined_tokens  = $token_counter->count_tokens( $combined_content );

				if ( $combined_tokens <= $max_tokens ) {
					// Merge: update pending with combined content
					$pending = array(
						'content'  => $combined_content,
						'source'   => 'content',
						'position' => $pending['position'],
						'length'   => strlen( $combined_content ),
					);
					continue;
				}
			}

			// Can't merge - save pending and move to next
			$merged[] = $pending;
			$pending  = $chunk;
		}

		// Don't forget the last pending chunk
		if ( null !== $pending ) {
			$merged[] = $pending;
		}

		return $merged;
	}

	/**
	 * Convert post HTML content to markdown format.
	 *
	 * Uses raw post_content without shortcode expansion to ensure embeddings
	 * remain stable (shortcode output can change without post updates).
	 *
	 * @param      \WP_Post $post The post to convert
	 * @return     string Markdown-formatted content
	 * @since     1.0.0
	 */
	public function convert_post_to_markdown( $post ) {
		$html = $post->post_content;
		if ( empty( trim( $html ) ) ) {
			return '';
		}

		// Wrap in minimal HTML structure to ensure proper parsing
		$html = '<html><body>' . $html . '</body></html>';

		try {
			$converter = new HtmlConverter(
				array(
					'strip_tags'      => true,
					'remove_nodes'    => 'script style',
					'hard_break'      => true,
					'header_style'    => 'atx',
					'suppress_errors' => true,
				)
			);

			// Enable table support (not enabled by default in league/html-to-markdown)
			$converter->getEnvironment()->addConverter( new TableConverter() );

			return $converter->convert( $html );
		} catch ( \Throwable $e ) {
			$error_message = String_Utils::replace_tags(
				__( 'Agenistrator: HTML to markdown conversion failed for post {post_id}: {error}', 'agenistrator' ),
				array(
					'post_id' => $post->ID,
					'error'   => $e->getMessage(),
				)
			);
			// phpcs:ignore WordPress.PHP.DevelopmentFunctions.error_log_error_log -- Production error logging for content extraction failures
			error_log( $error_message );

			// Return plain text fallback
			$plain_text = wp_strip_all_tags( $post->post_content );
			return trim( $plain_text );
		}
	}

	/**
	 * Reconstruct a chunk from a post based on stored metadata.
	 *
	 * Used to retrieve chunk content when serving search results without
	 * storing the actual chunk text in the database.
	 *
	 * @param      \WP_Post $post The post to extract chunk from
	 * @param      string $source Chunk source: 'title', 'excerpt', or 'content'
	 * @param      int|null $position Byte position in markdown content (for 'content' source)
	 * @param      int|null $length Byte length of chunk (for 'content' source)
	 * @param      string|null $expected_md5 Expected MD5 hash of rendered markdown (for 'content' source)
	 * @return     string|null The chunk content, or null if content has changed or invalid metadata
	 * @since     1.0.0
	 */
	public function reconstruct_chunk( $post, $source, $position = null, $length = null, $expected_md5 = null ) {
		switch ( $source ) {
			case 'title':
				return String_Utils::replace_tags(
					__( '{type}: {title}', 'agenistrator' ),
					array(
						'type'  => esc_html( ucfirst( $post->post_type ) ),
						'title' => esc_html( trim( $post->post_title ) ),
					)
				);

			case 'excerpt':
				return trim( $post->post_excerpt );

			case 'content':
				// Validate required parameters
				if ( null === $position || null === $length || null === $expected_md5 ) {
					return null;
				}

				// Re-render markdown and check if it has changed
				$markdown    = $this->convert_post_to_markdown( $post );
				$current_md5 = md5( $markdown );
				if ( $current_md5 !== $expected_md5 ) {
					return null;
				}

				// Extract chunk using byte offset/length
				return substr( $markdown, $position, $length );

			default:
				return null;
		}
	}

	/**
	 * Split text into chunks that fit within a token limit, tracking positions.
	 *
	 * First attempts to split at sentence boundaries, then falls back to
	 * word boundaries for very long sentences.
	 *
	 * @param      string $text The text to split
	 * @param      int $base_position Position of this text in the full markdown
	 * @param      int $max_tokens Maximum tokens per chunk
	 * @param      Token_Counter $token_counter Token counter instance
	 * @return     array Array of chunk objects with 'content', 'source', 'position', 'length'
	 */
	private function split_text_by_tokens_with_positions( $text, $base_position, $max_tokens, $token_counter ) {
		$chunks = array();

		// Try splitting at sentence boundaries first
		$sentences           = preg_split( '/(?<=[.!?])\s+/', $text, -1, PREG_SPLIT_NO_EMPTY );
		$current_chunk       = '';
		$current_chunk_start = 0;
		$current_pos_in_text = 0;

		foreach ( $sentences as $sentence ) {
			$sentence_trimmed = trim( $sentence );
			if ( empty( $sentence_trimmed ) ) {
				continue;
			}

			// Find position of this sentence in the text
			$sentence_pos = strpos( $text, $sentence_trimmed, $current_pos_in_text );
			if ( false === $sentence_pos ) {
				$sentence_pos = $current_pos_in_text;
			}

			$sentence_tokens = $token_counter->count_tokens( $sentence_trimmed );

			// If single sentence exceeds limit, split by words
			if ( $sentence_tokens > $max_tokens ) {
				// Flush current chunk first
				if ( ! empty( $current_chunk ) ) {
					$chunks[]      = array(
						'content'  => trim( $current_chunk ),
						'source'   => 'content',
						'position' => $base_position + $current_chunk_start,
						'length'   => strlen( trim( $current_chunk ) ),
					);
					$current_chunk = '';
				}
				// Split sentence by words with position tracking
				$word_chunks         = $this->split_by_words_with_positions(
					$sentence_trimmed,
					$base_position + $sentence_pos,
					$max_tokens,
					$token_counter
				);
				$chunks              = array_merge( $chunks, $word_chunks );
				$current_pos_in_text = $sentence_pos + strlen( $sentence_trimmed );
				$current_chunk_start = $current_pos_in_text;
				continue;
			}

			// Check if adding this sentence would exceed limit
			$test_chunk  = empty( $current_chunk ) ? $sentence_trimmed : $current_chunk . ' ' . $sentence_trimmed;
			$test_tokens = $token_counter->count_tokens( $test_chunk );

			if ( $test_tokens > $max_tokens ) {
				// Save current chunk and start new one
				if ( ! empty( $current_chunk ) ) {
					$chunks[] = array(
						'content'  => trim( $current_chunk ),
						'source'   => 'content',
						'position' => $base_position + $current_chunk_start,
						'length'   => strlen( trim( $current_chunk ) ),
					);
				}
				$current_chunk       = $sentence_trimmed;
				$current_chunk_start = $sentence_pos;
			} else {
				if ( empty( $current_chunk ) ) {
					$current_chunk_start = $sentence_pos;
				}
				$current_chunk = $test_chunk;
			}

			$current_pos_in_text = $sentence_pos + strlen( $sentence_trimmed );
		}

		// Don't forget the last chunk
		if ( ! empty( $current_chunk ) ) {
			$chunks[] = array(
				'content'  => trim( $current_chunk ),
				'source'   => 'content',
				'position' => $base_position + $current_chunk_start,
				'length'   => strlen( trim( $current_chunk ) ),
			);
		}

		return $chunks;
	}

	/**
	 * Split text by words when sentence splitting isn't sufficient, tracking positions.
	 *
	 * @param      string $text The text to split
	 * @param      int $base_position Position of this text in the full markdown
	 * @param      int $max_tokens Maximum tokens per chunk
	 * @param      Token_Counter $token_counter Token counter instance
	 * @return     array Array of chunk objects with 'content', 'source', 'position', 'length'
	 */
	private function split_by_words_with_positions( $text, $base_position, $max_tokens, $token_counter ) {
		$chunks = array();
		$words  = preg_split( '/\s+/', $text, -1, PREG_SPLIT_NO_EMPTY );

		$current_chunk       = '';
		$current_chunk_start = 0;
		$current_pos_in_text = 0;

		foreach ( $words as $word ) {
			// Find position of this word in the text
			$word_pos = strpos( $text, $word, $current_pos_in_text );
			if ( false === $word_pos ) {
				$word_pos = $current_pos_in_text;
			}

			$test_chunk  = empty( $current_chunk ) ? $word : $current_chunk . ' ' . $word;
			$test_tokens = $token_counter->count_tokens( $test_chunk );

			if ( $test_tokens > $max_tokens ) {
				if ( ! empty( $current_chunk ) ) {
					$chunks[] = array(
						'content'  => trim( $current_chunk ),
						'source'   => 'content',
						'position' => $base_position + $current_chunk_start,
						'length'   => strlen( trim( $current_chunk ) ),
					);
				}
				$current_chunk       = $word;
				$current_chunk_start = $word_pos;
			} else {
				if ( empty( $current_chunk ) ) {
					$current_chunk_start = $word_pos;
				}
				$current_chunk = $test_chunk;
			}

			$current_pos_in_text = $word_pos + strlen( $word );
		}

		if ( ! empty( $current_chunk ) ) {
			$chunks[] = array(
				'content'  => trim( $current_chunk ),
				'source'   => 'content',
				'position' => $base_position + $current_chunk_start,
				'length'   => strlen( trim( $current_chunk ) ),
			);
		}

		return $chunks;
	}
}
