<?php

namespace Limb_Chatbot\Includes\AI_Providers\Gemini\Endpoints\Generate_Content\Handlers;

use Limb_Chatbot\Includes\Chatbot_Tools\Actions\Action_Chatbot_Tool;
use Limb_Chatbot\Includes\Chatbot_Tools\Chatbot_Tools;
use Limb_Chatbot\Includes\Data_Objects\Action_Tool_Calls_Message;
use Limb_Chatbot\Includes\Data_Objects\Message;
use Limb_Chatbot\Includes\AI_Providers\Gemini\Endpoints\Generate_Content\Tool_Calls_Message;
use Limb_Chatbot\Includes\AI_Providers\Gemini\Handlers\Response_Handler;
use Limb_Chatbot\Includes\Data_Objects\Token_Usage;
use Limb_Chatbot\Includes\Exceptions\Error_Codes;
use Limb_Chatbot\Includes\Exceptions\Exception;

/**
 * Class Generate_Content_Response_Handler
 *
 * Handles the parsing and formatting of responses from Gemini's generateContent or streamGenerateContent endpoints.
 *
 * @package Limb_Chatbot\Includes\AI_Providers\Gemini\Endpoints\Generate_Content\Handlers
 * @since 1.0.0
 */
class Generate_Content_Response_Handler extends Response_Handler {

	/**
	 * The message object which should be returned.
	 *
	 * @var Message|Tool_Calls_Message|null
	 */
	protected $message;

	/**
	 * The generated content
	 *
	 * @var object|null
	 */
	protected $content;

	/**
	 * Returns the parsed message object.
	 *
	 * @return Message|Tool_Calls_Message|null
	 * @throws Exception
	 * @since 1.0.0
	 */
	public function get_message() {
		$this->set_message( $this->message_factory() );

		return $this->message;
	}

	/**
	 * Sets the message property.
	 *
	 * @param  Message|Tool_Calls_Message  $message  The message object to set.
	 *
	 * @return void
	 * @since 1.0.0
	 */
	public function set_message( $message ) {
		$this->message = $message;
	}

	/**
	 * Factory method to generate message based on response body (stream or non-stream).
	 *
	 * @return Message|Tool_Calls_Message|null
	 * @throws Exception If unsupported or empty response is encountered.
	 * @since 1.0.0
	 */
	public function message_factory() {
		if ( $this->is_stream ) {
			return $this->stream_message_factory();
		} else {
			$this->content = $this->get_content();
			foreach($this->content as $part) {
				if ( ! empty( $part->functionCall ) ) {
					$function_name = $part->functionCall->name;
					if ( Chatbot_Tools::instance()->get_tool( $function_name ) instanceof Action_Chatbot_Tool ) {
						return Action_Tool_Calls_Message::make( [
							'action_name' => $function_name,
							'role'        => Message::ROLE_ASSISTANT,
						] );
					}

					return Tool_Calls_Message::make( [
						'role'  => Message::ROLE_ASSISTANT,
						'parts' => $this->content,
					] );
				} elseif ( ! empty( $part->text ) ) {
					if ( ! empty( $this->get_body() ) && $usage = $this->get_body()->usageMetadata ) {
						$usage = Token_Usage::make( [
							'input_tokens'  => $usage->promptTokenCount,
							'output_tokens' => $usage->candidatesTokenCount,
						] );
					}

					return Message::make( [
						'role'    => Message::ROLE_ASSISTANT,
						'content' => [ [ 'type' => 'text', 'text' => [ 'value' => $part->text ] ] ],
						'usage'   => $usage ?? null
					] );
				}
			}
			throw new Exception(Error_Codes::GEMINI_ERROR, __('Unexpected response', 'limb-chatbot'));
		}
	}

	/**
	 * Builds a streamed message from chunked JSON responses.
	 *
	 * @return Message|Tool_Calls_Message
	 * @throws Exception If message parsing fails.
	 * @since 1.0.0
	 */
	private function stream_message_factory() {
		$content = '';
		$message = null;
		$function_call_parts = null;

		/**
		 * Below we have full chunks array. Which can contain any kind of choices - simple, tool-calls, etc ...
		 * Chunks with not empty content will be acted as single Message
		 * Chunks with tool-calls will be grouped into multiple Tool_Call_Messages - for later process them individually
		 */
		foreach ( $this->stream_parser->complete_jsons as $stream_chunk ) {
			$has_candidates = ! empty( $stream_chunk->candidates[0] );
			$has_content    = $has_candidates && ! empty( $stream_chunk->candidates[0]->content );
			$has_parts      = $has_content && isset( $stream_chunk->candidates[0]->content->parts );
			if ( ! $has_parts ) {
				continue;
			}
			$parts = $stream_chunk->candidates[0]->content->parts;
			if ( ! empty( $parts[0]->text ) ) {
				$content .= $parts[0]->text;
			} elseif ( ! empty( $parts[0]->functionCall ) ) {
				$function_call_parts = $parts;
				// Breaking here, cause there can't be any other iteration. functionCall is being passed as single chunk
				break;
			}
		}
		if ( ! empty( $function_call_parts ) ) {
			$function_name = $function_call_parts[0]->functionCall->name;
			if ( Chatbot_Tools::instance()->get_tool( $function_name ) instanceof Action_Chatbot_Tool ) {
				$message = Action_Tool_Calls_Message::make( [
					'action_name' => $function_name,
					'role'        => Message::ROLE_ASSISTANT,
				] );
			} else {
				$message = Tool_Calls_Message::make( [
					'role'  => Message::ROLE_ASSISTANT,
					'parts' => $function_call_parts
				] );
			}
		} elseif ( ! empty( $content ) ) {
			$chunk_with_usage = $this->stream_parser->complete_jsons[ count( $this->stream_parser->complete_jsons ) - 1 ];
			if ( $usage = $chunk_with_usage->usageMetadata ?? null ) {
				$usage = Token_Usage::make( [
					'input_tokens'  => $usage->promptTokenCount,
					'output_tokens' => $usage->candidatesTokenCount,
				] );
			}
			$message = Message::make( [
				'role'     => Message::ROLE_ASSISTANT,
				'content'  => [ [ 'type' => 'text', 'text' => [ 'value' => $content ] ] ],
				'streamed' => true,
				'usage'    => $usage ?? null
			] );
		}
		if ( empty( $message ) ) {
			if ( method_exists( $this->stream_parser, 'log_parser_state' ) ) {
				$this->stream_parser->log_parser_state();
			}
			throw new Exception( Error_Codes::TECHNICAL_ERROR, __( 'Error while parsing the streamed message', 'limb-chatbot' ) );
		}

		return $message;
	}

	/**
	 * Retrieve and format the AI response message.
	 *
	 * Extracts the first content candidate from the AI response, wraps it
	 * in a Message object, and optionally includes token usage metadata.
	 *
	 * @return Message  The formatted response message with role, content, and usage.
	 *
	 * @since 1.0.0
	 */
	public function get_response_message(){
		$content = ! empty( $this->get_body()->candidates[0]->content ) ? $this->get_body()->candidates[0]->content : null;
		if ( ! empty( $this->get_body() ) && $usage = $this->get_body()->usageMetadata ) {
			$usage = Token_Usage::make( [
				'input_tokens'  => $usage->promptTokenCount,
				'output_tokens' => $usage->candidatesTokenCount,
			] );
		}

		return Message::make( [
			'role'    => Message::ROLE_ASSISTANT,
			'content' => [ [ 'type' => 'text', 'text' => [ 'value' => $content->parts[0]->text ] ] ],
			'usage'   => $usage ?? null
		] );
	}

	/**
	 * Decode JSON content from the AI response safely.
	 *
	 * Extracts the first candidate's text, removes Markdown-style code blocks (```json ... ```),
	 * and decodes it into an associative array.
	 *
	 * @return array Decoded JSON data or empty array if invalid.
	 * @since 1.0.0
	 */
	public function get_json() {
		$text            = ! empty( $this->get_body()->candidates[0]->content->parts[0]->text ) ? $this->get_body()->candidates[0]->content->parts[0]->text : null;
		if ( ! empty( $text ) ) {
			// If the content comes with ```json ... ``` code block
			$text = preg_replace( '/^```(?:json)?\s*|\s*```$/', '', trim( $text ) );
			$text = trim( $text );
			// Decode the JSON safely
			$entities = json_decode( $text, true );
			if ( json_last_error() !== JSON_ERROR_NONE ) {
				return [];
			}
			return $entities;
		}

		return [];
	}

	private function get_content() {
		$this->content = ! empty( $this->get_body()->candidates[0]->content ) ? $this->get_body()->candidates[0]->content->parts : null;
		if ( ! empty( $this->content ) && is_array( $this->content ) ) {
			usort( $this->content, static function ( $a, $b ) {
				$a_has_function = ! empty( $a->functionCall );
				$b_has_function = ! empty( $b->functionCall );

				// Move items with functionCall to the front
				if ( $a_has_function === $b_has_function ) {
					return 0; // same type, keep original order
				}

				return $a_has_function ? - 1 : 1; // functionCall first
			} );
		}
		return $this->content;
	}
}