<?php

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

/**
 * Class Chat_Completion_Response_Handler
 *
 * Handles parsing and message extraction for OpenRouter chat completions (including streamed responses).
 *
 * @since 1.0.14
 */
class Chat_Completion_Response_Handler extends Response_Handler {

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

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

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

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

	/**
	 * Returns the response message for non-streamed responses.
	 *
	 * @return Message|null
	 * @throws Exception
	 * @since 1.0.14
	 */
	public function get_response_message(): ?Message {
		$body = $this->get_body();
		if ( ! empty( $body->choices[0]->message->content ) ) {
			return Message::make( [
				'role'    => $body->choices[0]->message->role,
				'content' => [
					[
						'type' => 'text',
						'text' => [ 'value' => $body->choices[0]->message->content ]
					]
				]
			] );
		}

		return null;
	}

	/**
	 * Returns JSON-decoded response body.
	 *
	 * @return array|null
	 * @since 1.0.14
	 */
	public function get_json(): ?array {
		$body = $this->get_body();
		if ( ! empty( $body->choices[0]->message->content ) ) {
			return json_decode( $body->choices[0]->message->content, true );
		}

		return null;
	}

	/**
	 * 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.14
	 */
	private function message_factory() {
		if ( $this->is_stream ) {
			return $this->stream_message_factory();
		} else {
			$message = ! empty( $this->get_body()->choices[0]->message ) ? $this->get_body()->choices[0]->message : null;
			if ( ! empty( $message->tool_calls ) ) {
				if ( ! empty( $message->tool_calls[0]->function->name ) && Chatbot_Tools::instance()->get_tool( $message->tool_calls[0]->function->name ) instanceof Action_Chatbot_Tool ) {
					return Action_Tool_Calls_Message::make( [
						'action_name' => $message->tool_calls[0]->function->name,
						'role'        => Message::ROLE_ASSISTANT,
						'parts'       => $message->tool_calls,
					] );
				}
				return Tool_Calls_Message::make( [ 'role' => $message->role, 'tool_calls' => $message->tool_calls ] );
			} elseif ( ! empty( $message->content ) ) {
				return Message::make( [
					'role'    => $message->role,
					'content' => [
						[
							'type' => 'text',
							'text' => [ 'value' => $message->content ]
						]
					],
					'usage'   => Token_Usage::make( [
						'input_tokens'  => $this->get_body()->usage->prompt_tokens ?? 0,
						'output_tokens' => $this->get_body()->usage->completion_tokens ?? 0,
					] )
				] );
			}
		}

		return null;
	}

	/**
	 * Constructs a streamed message by merging multiple streamed chunks.
	 *
	 * @return Message|Tool_Calls_Message
	 * @throws Exception If stream data is invalid or no valid message can be built.
	 * @since 1.0.14
	 */
	private function stream_message_factory() {
		$content    = '';
		$tool_calls = [];

		foreach ( $this->stream_chunks as $stream_chunk ) {
			if ( empty( $stream_chunk->choices ) ) {
				continue;
			}
			$delta = $stream_chunk->choices[0]->delta;
			if ( ! empty( $delta->content ) ) {
				$content .= $delta->content;
			} elseif ( ! empty( $delta->tool_calls[0] ) ) {
				$tool_call    = $delta->tool_calls[0];
				$active_index = $tool_call->index;
				if ( ! isset( $tool_calls[ $active_index ] ) ) {
					$tool_calls[ $active_index ] = $tool_call;
					continue;
				}
				if ( isset( $tool_calls[ $active_index ]->{$tool_calls[ $active_index ]->type}->arguments ) ) {
					$tool_calls[ $active_index ]->{$tool_calls[ $active_index ]->type}->arguments .= $tool_call->{$tool_calls[ $active_index ]->type}->arguments;
				}
			}
		}

		$chunk_with_usage = $this->stream_chunks[ count( $this->stream_chunks ) - 1 ] ?? null;
		$usage = null;
		if ( ! empty( $chunk_with_usage->usage ) ) {
			$usage = Token_Usage::make( [
				'input_tokens'  => $chunk_with_usage->usage->prompt_tokens ?? 0,
				'output_tokens' => $chunk_with_usage->usage->completion_tokens ?? 0,
			] );
		}

		if ( ! empty( $tool_calls ) ) {
			$tool_calls = array_values( $tool_calls );
			if ( ! empty( $tool_calls[0]->function->name ) && Chatbot_Tools::instance()->get_tool( $tool_calls[0]->function->name ) instanceof Action_Chatbot_Tool ) {
				return Action_Tool_Calls_Message::make( [
					'action_name' => $tool_calls[0]->function->name,
					'role'        => Message::ROLE_ASSISTANT,
					'parts'       => $tool_calls,
				] );
			}
			return Tool_Calls_Message::make( [ 'role' => Message::ROLE_ASSISTANT, 'tool_calls' => $tool_calls ] );
		}

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