<?php

namespace Limb_Chatbot\Includes\AI_Providers\Open_Router\Services;

use Limb_Chatbot\Includes\Data_Objects\Chatbot;
use Limb_Chatbot\Includes\Data_Objects\Message;
use Limb_Chatbot\Includes\Exceptions\Error_Codes;
use Limb_Chatbot\Includes\Exceptions\Exception;
use Limb_Chatbot\Includes\Interfaces\Token_Calculator_Interface;

/**
 * Class Token_Calculator
 *
 * Estimates the number of tokens in a given message based on its text or attachments.
 * This is useful for validating against model limits and usage policies in AI interactions.
 *
 * @package Limb_Chatbot\Includes\AI_Providers\Open_Router\Services
 * @since 1.0.14
 */
class Token_Calculator implements Token_Calculator_Interface {

	const SPECIAL_CHARACTERS_PATTERN = '/[\x{1F600}-\x{1F64F}\x{1F300}-\x{1F5FF}\x{1F680}-\x{1F6FF}\x{1F700}-\x{1F77F}\x{1F780}-\x{1F7FF}\x{1F800}-\x{1F8FF}\x{1F900}-\x{1F9FF}\x{1FA00}-\x{1FA6F}\x{1FA70}-\x{1FAFF}\x{2600}-\x{26FF}\x{2700}-\x{27BF}]/u';
	const SMALL_WORD_THRESHOLD       = 4;
	const MEDIUM_WORD_THRESHOLD      = 8;
	const CHAR_PER_TOKEN             = 4;
	const ADDITIONAL_TOKEN_FOR_SPECIAL_CHARS = 1;

	const MAX_DIMENSION = 2048;
	const MIN_DIMENSION = 768;
	const TILE_SIZE     = 512;
	const TILE_TOKENS   = 170;
	const BASE_TOKENS   = 85;

	/**
	 * The message instance for which token calculation is performed.
	 *
	 * @var Message
	 */
	protected Message $message;

	/**
	 * Token_Calculator constructor.
	 *
	 * @param Message $message Message object containing content to be tokenized.
	 * @param Chatbot|null $chatbot Optional chatbot context, currently unused.
	 * @since 1.0.14
	 */
	public function __construct( Message $message, ?Chatbot $chatbot = null ) {
		$this->message = $message;
	}

	/**
	 * Calculates the estimated total token usage for the given message.
	 *
	 * Supports both text and base64-encoded image attachments.
	 *
	 * @return int Total token count.
	 * @throws Exception If attachment data is invalid.
	 * @since 1.0.14
	 */
	public function calculate(): int {
		$tokens = 0;
		foreach ( $this->message->get_content() as $content ) {
			$method = "count_{$content['type']}_tokens";
			if ( method_exists( $this, $method ) ) {
				$tokens += $this->$method( $content[ $content['type'] ]['value'] );
			}
		}

		return $tokens;
	}

	/**
	 * Calculates token count for a plain text input.
	 *
	 * @param string $text Plain text content.
	 * @return int Token estimate.
	 * @since 1.0.14
	 */
	private function count_text_tokens( $text ): int {
		$text        = trim( preg_replace( '/\s+/', ' ', $text ) );
		$words       = preg_split( '/(\W)/u', $text, -1, PREG_SPLIT_DELIM_CAPTURE | PREG_SPLIT_NO_EMPTY );
		$token_count = 0;
		foreach ( $words as $word ) {
			if ( trim( $word ) !== '' ) {
				if ( mb_strlen( $word ) < self::SMALL_WORD_THRESHOLD ) {
					$token_count += 1;
				} elseif ( mb_strlen( $word ) <= self::MEDIUM_WORD_THRESHOLD ) {
					$token_count += 2;
				} else {
					$token_count += ceil( mb_strlen( $word ) / self::CHAR_PER_TOKEN );
				}
				if ( preg_match( self::SPECIAL_CHARACTERS_PATTERN, $word ) ) {
					$token_count += self::ADDITIONAL_TOKEN_FOR_SPECIAL_CHARS;
				}
			}
		}

		return $token_count;
	}

	/**
	 * Calculates tokens for a base64-encoded image (vision API).
	 *
	 * @param string $base64 Base64-encoded image string.
	 * @return int Token estimate for the image.
	 * @throws Exception If the image data is invalid.
	 * @since 1.0.14
	 */
	private function count_attachment_tokens( $base64 ): int {
		$image = @imagecreatefromstring( base64_decode( preg_replace( '#^data:image/\w+;base64,#i', '', $base64 ) ) );
		if ( ! $image ) {
			throw new Exception( Error_Codes::DATA_INVALID_ATTACHMENT, __( 'Invalid image attachment', 'limb-chatbot' ) );
		}
		$width  = imagesx( $image );
		$height = imagesy( $image );
		imagedestroy( $image );

		if ( $width > self::MAX_DIMENSION || $height > self::MAX_DIMENSION ) {
			$scale  = min( self::MAX_DIMENSION / $width, self::MAX_DIMENSION / $height );
			$width  = (int) ( $width * $scale );
			$height = (int) ( $height * $scale );
		}

		$shortest_side = min( $width, $height );
		$scale         = max( 1, self::MIN_DIMENSION / $shortest_side );
		$scaled_width  = (int) min( $width * $scale, self::MAX_DIMENSION );
		$scaled_height = (int) min( $height * $scale, self::MAX_DIMENSION );

		$tiles_x = ceil( $scaled_width / self::TILE_SIZE );
		$tiles_y = ceil( $scaled_height / self::TILE_SIZE );

		return self::BASE_TOKENS + ( self::TILE_TOKENS * $tiles_x * $tiles_y );
	}
}
