<?php

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

use Limb_Chatbot\Includes\AI_Providers\Claude\Endpoints\Chat_Completion\Tool_Calls_Message;
use Limb_Chatbot\Includes\AI_Providers\Claude\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\Dataset_Entry;
use Limb_Chatbot\Includes\Data_Objects\Message;
use Limb_Chatbot\Includes\Data_Objects\Token_Usage;
use Limb_Chatbot\Includes\Exceptions\Error_Codes;
use Limb_Chatbot\Includes\Exceptions\Exception;
use Limb_Chatbot\Includes\Services\Collection;
use Limb_Chatbot\Includes\Services\Helper;
use WP_Http;

/**
 * Class Response_Handler
 *
 * Handles parsing and message extraction for Claude Messages API responses.
 *
 * @package Limb_Chatbot\Includes\AI_Providers\Claude\Endpoints\Chat_Completion
 * @since 1.0.9
 */
class Chat_Completion_Response_Handler extends Response_Handler {

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

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

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

	/**
	 * Token usage information extracted from the response.
	 *
	 * @var Token_Usage|null
	 * @since 1.0.9
	 */
	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.9
	 */
	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.9
	 */
	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.9
	 */
	public function get_message() {
		return $this->claude_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.9
	 */
	private function claude_message_factory() {
		if ( $this->is_stream ) {
			return $this->claude_stream_message_factory();
		}

		$body = $this->get_body();

		if ( empty( $body ) || empty( $body->content ) ) {
			return null;
		}

		$content = $body->content;

		// Check for tool_use blocks
		$tool_uses = array_values( array_filter( (array) $content,
			fn( $block ) => ( $block->type ?? '' ) === 'tool_use' ) );

		if ( ! empty( $tool_uses ) ) {
			$tool_uses_array = array_values( array_map( function ( $block ) {
				return array(
					'id'    => $block->id,
					'name'  => $block->name,
					'input' => (array) $block->input,
				);
			}, $tool_uses ) );

			// Check if this is an action tool call
			$first_tool_name = $tool_uses_array[0]['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_uses_array ),
				) );
			}

			$message = Tool_Calls_Message::make( array(
				'role' => Message::ROLE_ASSISTANT,
			) );
			$message->set_tool_uses( $tool_uses_array );
			$message->set_content_blocks( array_map( fn( $b ) => (array) $b, (array) $content ) );

			// Set usage
			if ( ! empty( $body->usage ) ) {
				$message->set_usage( Token_Usage::make( array(
					'input_tokens'  => $body->usage->input_tokens ?? 0,
					'output_tokens' => $body->usage->output_tokens ?? 0,
				) ) );
			}

			return $message;
		}

		// Extract text content
		$text_content = '';
		foreach ( $content as $block ) {
			if ( ( $block->type ?? '' ) === 'text' ) {
				$text_content .= $block->text;
			}
		}

		if ( ! empty( $text_content ) ) {
			return Message::make( array(
				'role'    => Message::ROLE_ASSISTANT,
				'content' => array(
					array(
						'type' => 'text',
						'text' => array( 'value' => $text_content ),
					),
				),
				'usage'   => ! empty( $body->usage ) ? Token_Usage::make( array(
					'input_tokens'  => $body->usage->input_tokens ?? 0,
					'output_tokens' => $body->usage->output_tokens ?? 0,
				) ) : null,
			) );
		}

		return null;
	}

	/**
	 * Normalize tool calls for action tool format.
	 *
	 * @param  array  $tool_uses  Array of tool use blocks.
	 *
	 * @return array Normalized for Action_Tool_Calls_Message.
	 * @since 1.0.9
	 */
	private function normalize_tool_calls_for_action( array $tool_uses ): array {
		return array_map( function ( $tool_use ) {
			return (object) array(
				'id'       => $tool_use['id'],
				'type'     => 'function',
				'function' => (object) array(
					'name'      => $tool_use['name'],
					'arguments' => wp_json_encode( $tool_use['input'] ),
				),
			);
		}, $tool_uses );
	}

	/**
	 * Constructs a streamed message by merging multiple streamed chunks.
	 *
	 * @return Message|Tool_Calls_Message
	 * @throws Exception If stream data is invalid.
	 * @since 1.0.9
	 */
	private function claude_stream_message_factory() {
		$content          = '';
		$tool_uses        = array();
		$usage            = null;
		$current_tool_use = null;

		foreach ( $this->stream_chunks as $chunk ) {
			$type = $chunk->type ?? '';

			switch ( $type ) {
				case 'content_block_start':
					if ( isset( $chunk->content_block->type ) && $chunk->content_block->type === 'tool_use' ) {
						$current_tool_use = array(
							'id'    => $chunk->content_block->id ?? '',
							'name'  => $chunk->content_block->name ?? '',
							'input' => array(),
						);
					}
					break;

				case 'content_block_delta':
					$delta = $chunk->delta ?? null;
					if ( $delta ) {
						if ( ( $delta->type ?? '' ) === 'text_delta' ) {
							$content .= $delta->text ?? '';
						} elseif ( ( $delta->type ?? '' ) === 'input_json_delta' && $current_tool_use !== null ) {
							if ( ! isset( $current_tool_use['partial_json'] ) ) {
								$current_tool_use['partial_json'] = '';
							}
							$current_tool_use['partial_json'] .= $delta->partial_json ?? '';
						}
					}
					break;

				case 'content_block_stop':
					if ( $current_tool_use !== null ) {
						if ( ! empty( $current_tool_use['partial_json'] ) ) {
							$parsed = json_decode( $current_tool_use['partial_json'], true );
							if ( $parsed !== null ) {
								$current_tool_use['input'] = $parsed;
							}
							unset( $current_tool_use['partial_json'] );
						}
						$tool_uses[]      = $current_tool_use;
						$current_tool_use = null;
					}
					break;

				case 'message_delta':
					if ( isset( $chunk->usage ) ) {
						$usage = Token_Usage::make( array(
							'input_tokens'  => $chunk->usage->input_tokens ?? 0,
							'output_tokens' => $chunk->usage->output_tokens ?? 0,
						) );
					}
					break;

				case 'message_start':
					if ( isset( $chunk->message->usage ) ) {
						$usage = Token_Usage::make( array(
							'input_tokens'  => $chunk->message->usage->input_tokens ?? 0,
							'output_tokens' => $chunk->message->usage->output_tokens ?? 0,
						) );
					}
					break;
			}
		}

		if ( ! empty( $tool_uses ) ) {
			$tool_name = $tool_uses[0]['name'] ?? '';
			$tool      = ! empty( $tool_name ) ? Chatbot_Tools::instance()->get_tool( $tool_name ) : null;

			if ( $tool instanceof Action_Chatbot_Tool ) {
				$action_message = Action_Tool_Calls_Message::make( array(
					'action_name' => $tool_name,
					'role'        => Message::ROLE_ASSISTANT,
					'parts'       => $this->normalize_tool_calls_for_action( $tool_uses ),
				) );

				if ( $usage ) {
					$action_message->set_usage( $usage );
				}

				return $action_message;
			}

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

			if ( $usage ) {
				$message->set_usage( $usage );
			}

			return $message;
		}

		if ( ! empty( $content ) ) {
			return Message::make( array(
				'role'     => Message::ROLE_ASSISTANT,
				'content'  => array( array( 'type' => 'text', 'text' => array( 'value' => $content ) ) ),
				'usage'    => $usage,
				'streamed' => true,
			) );
		}

		// Use attach_utility for consistent error message format with non-streaming
		$exception = new Exception(
			Error_Codes::TECHNICAL_ERROR,
			__( 'Error while parsing the streamed message', 'limb-chatbot' )
		);
		throw $exception->attach_utility( $this->get_endpoint()->get_utility(), true );
	}

	/**
	 * Get the response as a simple Message object.
	 *
	 * @return Message|null
	 * @throws Exception If parsing fails.
	 * @since 1.0.9
	 */
	public function get_response_message(): ?Message {
		$message = $this->get_message();
		if ( $message instanceof Message ) {
			return $message;
		}

		return null;
	}

	/**
	 * Get dataset entries from the response.
	 *
	 * @return Collection
	 * @since 1.0.9
	 */
	public function get_dataset_entries(): Collection {
		$dataset_entries = new Collection();
		$body            = $this->get_body();

		if ( empty( $body->content ) ) {
			return $dataset_entries;
		}

		$text_content = '';
		foreach ( $body->content as $block ) {
			if ( ( $block->type ?? '' ) === 'text' ) {
				$text_content .= $block->text;
			}
		}

		if ( ! empty( $text_content ) ) {
			$content = preg_replace( '/^```(?:json)?\s*|\s*```$/', '', trim( $text_content ) );
			if ( Helper::is_probable_json( $content ) ) {
				$entry = json_decode( $content );
				if ( $entry ) {
					$dataset_entry        = Dataset_Entry::make();
					$dataset_entry->entry = array(
						'system'         => '',
						'messages'       => array(
							array(
								'input'  => $entry->question ?? '',
								'output' => $entry->answer ?? '',
							),
						),
						'entry_metadata' => array(
							'entry_type' => 'qa',
						),
					);
					$dataset_entries->push_item( $dataset_entry );
				}
			}
		}

		return $dataset_entries;
	}

	/**
	 * Get the JSON-decoded content from the response.
	 *
	 * @return array
	 * @since 1.0.9
	 */
	public function get_json(): array {
		$body = $this->get_body();

		if ( empty( $body->content ) ) {
			return array();
		}

		$text_content = '';
		foreach ( $body->content as $block ) {
			if ( ( $block->type ?? '' ) === 'text' ) {
				$text_content .= $block->text;
			}
		}

		if ( ! empty( $text_content ) ) {
			$content = preg_replace( '/^```(?:json)?\s*|\s*```$/', '', trim( $text_content ) );
			$decoded = json_decode( $content, true );
			if ( $decoded !== null ) {
				return $decoded;
			}
		}

		return array();
	}
}

