<?php

namespace Limb_Chatbot\Includes\AI_Providers\Grok\Endpoints\Chat_Completion\Handlers;

use Limb_Chatbot\Includes\AI_Providers\Grok\Endpoints\Chat_Completion\Tool_Calls_Message;
use Limb_Chatbot\Includes\AI_Providers\Grok\Handlers\Response_Handler;
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\Data_Objects\Token_Usage;
use Limb_Chatbot\Includes\Exceptions\Exception;
use WP_Http;

/**
 * Class Chat_Completion_Response_Handler
 *
 * Handles parsing and message extraction for xAI Chat Completions API responses.
 *
 * @package Limb_Chatbot\Includes\AI_Providers\Grok\Endpoints\Chat_Completion
 * @since 1.0.12
 */
class Chat_Completion_Response_Handler extends Response_Handler {

	/**
	 * When true, response bodies will be decoded as associative arrays.
	 *
	 * @var bool|null
	 * @since 1.0.12
	 */
	protected ?bool $body_cast_array = null;

	/**
	 * The parsed message object returned after handling the response.
	 *
	 * @var Message|Tool_Calls_Message|null
	 * @since 1.0.12
	 */
	protected $message;

	/**
	 * Parsed stream chunks (decoded from JSON) for streamed responses.
	 *
	 * @var array|null
	 * @since 1.0.12
	 */
	protected $stream_chunks;

	/**
	 * Token usage information extracted from the response.
	 *
	 * @var Token_Usage|null
	 * @since 1.0.12
	 */
	protected ?Token_Usage $usage = null;

	/**
	 * Parses the raw response from the API.
	 *
	 * For streamed responses, it collects and decodes all complete JSON chunks.
	 *
	 * @since 1.0.12
	 */
	public function parse() {
		if ( $this->is_stream ) {
			$this->stream_chunks = $this->stream_parser->complete_jsons ?? array();
		}
		parent::parse();
	}

	/**
	 * Defines and decodes the response body.
	 *
	 * @return void
	 * @since 1.0.12
	 */
	public function define_body(): void {
		if ( get_class( $this->http_client ) === WP_Http::class ) {
			$body       = wp_remote_retrieve_body( $this->response );
			$this->body = $this->is_application_json() ? json_decode( $body, $this->body_cast_array ) : $body;
		} else {
			$this->body = null;
		}
	}

	/**
	 * Returns the parsed message object (text or tool-call based).
	 *
	 * @return Message|Tool_Calls_Message|null
	 * @throws Exception If parsing fails.
	 * @since 1.0.12
	 */
	public function get_message() {
		return $this->message_factory();
	}

	/**
	 * Determines how to construct the message depending on stream state.
	 *
	 * @return Message|Tool_Calls_Message|null
	 * @throws Exception If parsing fails or message structure is unknown.
	 * @since 1.0.12
	 */
	private function message_factory() {
		if ( $this->is_stream ) {
			return $this->stream_message_factory();
		}

		$body = $this->get_body();

		if ( empty( $body ) || empty( $body->choices[0] ) ) {
			return null;
		}

		$choice = $body->choices[0];

		// Check for tool calls (OpenAI-compatible format)
		if ( isset( $choice->message->tool_calls ) && ! empty( $choice->message->tool_calls ) ) {
			$tool_calls = $choice->message->tool_calls;

			// Check if this is an action tool call
			$first_tool_name = $tool_calls[0]->function->name ?? '';
			$first_tool      = ! empty( $first_tool_name ) ? Chatbot_Tools::instance()->get_tool( $first_tool_name ) : null;

			if ( $first_tool instanceof Action_Chatbot_Tool ) {
				return Action_Tool_Calls_Message::make( array(
					'action_name' => $first_tool_name,
					'role'        => Message::ROLE_ASSISTANT,
					'parts'       => $this->normalize_tool_calls_for_action( $tool_calls ),
				) );
			}

			$message = Tool_Calls_Message::make( array(
				'role' => Message::ROLE_ASSISTANT,
			) );
			$message->set_tool_calls( $tool_calls );

			return $message;
		}

		// Regular text response
		$content = $choice->message->content ?? '';

		return Message::make( array(
			'role'    => Message::ROLE_ASSISTANT,
			'content' => array(
				array(
					'type' => 'text',
					'text' => array( 'value' => $content ),
				),
			),
		) );
	}

	/**
	 * Creates a message from streamed response chunks.
	 *
	 * @return Message|Tool_Calls_Message|null
	 * @since 1.0.12
	 */
	private function stream_message_factory() {
		if ( empty( $this->stream_parser ) ) {
			return null;
		}

		$stop_reason = $this->stream_parser->get_stop_reason();

		// Check for tool calls
		if ( $stop_reason === 'tool_calls' ) {
			$tool_calls = $this->stream_parser->get_tool_calls();

			if ( ! empty( $tool_calls ) ) {
				$first_tool_name = $tool_calls[0]['function']['name'] ?? '';
				$first_tool      = ! empty( $first_tool_name ) ? Chatbot_Tools::instance()->get_tool( $first_tool_name ) : null;

				if ( $first_tool instanceof Action_Chatbot_Tool ) {
					return Action_Tool_Calls_Message::make( array(
						'action_name' => $first_tool_name,
						'role'        => Message::ROLE_ASSISTANT,
						'parts'       => $this->normalize_stream_tool_calls_for_action( $tool_calls ),
					) );
				}

				$message = Tool_Calls_Message::make( array(
					'role' => Message::ROLE_ASSISTANT,
				) );
				$message->set_tool_calls_from_array( $tool_calls );

				return $message;
			}
		}

		// Regular text response
		$content = $this->stream_parser->get_text_content();

		return Message::make( array(
			'role'    => Message::ROLE_ASSISTANT,
			'content' => array(
				array(
					'type' => 'text',
					'text' => array( 'value' => $content ),
				),
			),
		) );
	}

	/**
	 * Normalize tool calls for action processing.
	 *
	 * @param  array  $tool_calls  Tool calls from API response.
	 *
	 * @return array Normalized tool calls.
	 * @since 1.0.12
	 */
	private function normalize_tool_calls_for_action( $tool_calls ): array {
		$normalized = array();

		foreach ( $tool_calls as $tool_call ) {
			$normalized[] = array(
				'id'        => $tool_call->id,
				'name'      => $tool_call->function->name,
				'arguments' => json_decode( $tool_call->function->arguments, true ) ?? array(),
			);
		}

		return $normalized;
	}

	/**
	 * Normalize streamed tool calls for action processing.
	 *
	 * @param  array  $tool_calls  Tool calls from stream parser.
	 *
	 * @return array Normalized tool calls.
	 * @since 1.0.12
	 */
	private function normalize_stream_tool_calls_for_action( array $tool_calls ): array {
		$normalized = array();

		foreach ( $tool_calls as $tool_call ) {
			$normalized[] = array(
				'id'        => $tool_call['id'],
				'name'      => $tool_call['function']['name'],
				'arguments' => json_decode( $tool_call['function']['arguments'], true ) ?? array(),
			);
		}

		return $normalized;
	}

	/**
	 * Get usage information from the response.
	 *
	 * @return Token_Usage|null
	 * @since 1.0.12
	 */
	public function get_usage(): ?Token_Usage {
		if ( $this->usage ) {
			return $this->usage;
		}

		$usage_data = null;

		if ( $this->is_stream && $this->stream_parser ) {
			$usage_data = $this->stream_parser->get_usage();
		} elseif ( isset( $this->get_body()->usage ) ) {
			$usage_data = $this->get_body()->usage;
		}

		if ( $usage_data ) {
			$this->usage = Token_Usage::make( array(
				'prompt_tokens'     => $usage_data->prompt_tokens ?? 0,
				'completion_tokens' => $usage_data->completion_tokens ?? 0,
				'total_tokens'      => $usage_data->total_tokens ?? 0,
			) );
		}

		return $this->usage;
	}
}
