<?php

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

use Limb_Chatbot\Includes\AI_Providers\Claude\Endpoints\Chat_Completion\Handlers\Chat_Completion_Response_Handler;
use Limb_Chatbot\Includes\AI_Providers\Claude\Endpoints\Claude_Endpoint;
use Limb_Chatbot\Includes\AI_Providers\Claude\Services\Message_Service;
use Limb_Chatbot\Includes\AI_Providers\Claude\Utilities\Chatbot_Utility;
use Limb_Chatbot\Includes\AI_Providers\Claude\Utilities\Copilot_Utility;
use Limb_Chatbot\Includes\Data_Objects\AI_Model;
use Limb_Chatbot\Includes\Data_Objects\Message;
use Limb_Chatbot\Includes\Exceptions\Error_Codes;
use Limb_Chatbot\Includes\Exceptions\Exception;
use Limb_Chatbot\Includes\Services\Collection;

/**
 * Class Chat_Completion_Endpoint
 *
 * Handles Claude Messages API interactions for assistant responses.
 *
 * @package Limb_Chatbot\Includes\AI_Providers\Claude\Endpoints\Chat_Completion
 * @since 1.0.9
 */
class Chat_Completion_Endpoint extends Claude_Endpoint {

	/**
	 * Default max tokens if not specified.
	 *
	 * @var int
	 * @since 1.0.9
	 */
	const DEFAULT_MAX_TOKENS = 4096;

	/**
	 * Chat_Completion_Endpoint constructor.
	 *
	 * @param  Chatbot_Utility|Copilot_Utility  $utility  Utility instance for configuration.
	 *
	 * @since 1.0.9
	 */
	public function __construct( $utility ) {
		parent::__construct( $utility );
	}

	/**
	 * Sends a messages request to the Claude API.
	 *
	 * Prepares the message body, attaches tools and streaming configs,
	 * and returns the AI-generated message.
	 *
	 * @return Message|Tool_Calls_Message The generated message from Claude.
	 * @throws Exception If messages are empty or invalid.
	 * @since 1.0.9
	 */
	public function generate() {
		$model = $this->utility->global_utility->get_ai_model();
		$is_stream   = $this->utility->global_utility->get_stream();


		if ( ! $model ) {
			throw new Exception( Error_Codes::AI_MODEL_NOT_SET, __( 'AI Model is not set.', 'limb-chatbot' ) );
		}

		$messages    = $this->prepare_messages( $model );
		$http_client = $this->http_client_factory( $is_stream );

		$body = array(
			'model'      => $model->get_name(),
			'max_tokens' => $this->get_max_tokens(),
			'messages'   => $messages['messages'],
		);

		// Add system prompt if present
		if ( ! empty( $messages['system'] ) ) {
			$body['system'] = $messages['system'];
		}

		// Add tools if present
		$tools = $this->utility->get_tools();
		if ( ! empty( $tools ) ) {
			$body['tools'] = $tools;
		}

		// Set temperature
		$body['temperature'] = 0; // Make this dynamic

		// Enable streaming
		if ( $is_stream ) {
			$body['stream'] = true;
		}

		$args = array(
			'body'    => wp_json_encode( $body ),
			'headers' => $this->get_header(),
			'timeout' => $this->utility->get_timeout() ?? 120,
		);

		// Configure stream parser for streaming responses
		if ( $is_stream ) {
			$stream_parser         = new Stream_Parser();
			$args['stream_parser'] = array( $stream_parser, 'parser' );
		}

		$response = $http_client->post( self::API_BASE_URL . '/messages', $args );


		return ( new Chat_Completion_Response_Handler( $response, $http_client, $this, $stream_parser ?? null ) )->get_message();
	}

	/**
	 * Prepares an array of messages for the API request.
	 *
	 * Extracts system prompt and formats messages for Claude.
	 *
	 * @param  AI_Model  $model  The model for message formatting.
	 *
	 * @return array Array with 'system' and 'messages' keys.
	 * @throws Exception If the message list is empty.
	 * @since 1.0.9
	 */
	private function prepare_messages( $model ) {
		$messages       = array();
		$system_content = null;

		foreach ( $this->utility->global_utility->get_messages() as $message ) {
			$prepared = $this->prepare_individual_message( $message, $model );
			if ( $prepared === null ) {
				continue;
			}

			// Extract system prompt (Claude uses top-level system parameter)
			if ( isset( $prepared['is_system'] ) && $prepared['is_system'] ) {
				$system_content = $prepared['content'];
				continue;
			}

			$messages[] = $prepared;
		}

		if ( empty( $messages ) ) {
			throw new Exception( Error_Codes::CHATBOT_NO_MESSAGES, __( 'Messages are empty', 'limb-chatbot' ) );
		}

		// Merge consecutive messages with the same role
		$messages = $this->merge_consecutive_messages( $messages );

		return array(
			'system'   => $system_content,
			'messages' => $messages,
		);
	}

	/**
	 * Merges consecutive messages with the same role.
	 *
	 * Claude requires alternating user/assistant turns.
	 *
	 * @param  array  $messages  The array of messages.
	 *
	 * @return array Merged messages array.
	 * @since 1.0.9
	 */
	private function merge_consecutive_messages( array $messages ): array {
		if ( empty( $messages ) ) {
			return $messages;
		}

		$merged    = array();
		$last_role = null;
		$last_idx  = - 1;

		foreach ( $messages as $message ) {
			if ( $message['role'] === $last_role && $last_idx >= 0 ) {
				// Merge content
				$existing_content = $merged[ $last_idx ]['content'];
				$new_content      = $message['content'];

				if ( is_string( $existing_content ) && is_string( $new_content ) ) {
					$merged[ $last_idx ]['content'] = $existing_content . "\n\n" . $new_content;
				} elseif ( is_array( $existing_content ) && is_array( $new_content ) ) {
					$merged[ $last_idx ]['content'] = array_merge( $existing_content, $new_content );
				} elseif ( is_string( $existing_content ) ) {
					$merged[ $last_idx ]['content'] = array_merge(
						array( array( 'type' => 'text', 'text' => $existing_content ) ),
						is_array( $new_content ) ? $new_content : array(
							array(
								'type' => 'text',
								'text' => $new_content
							)
						)
					);
				} else {
					$merged[ $last_idx ]['content'] = array_merge(
						$existing_content,
						is_array( $new_content ) ? $new_content : array(
							array(
								'type' => 'text',
								'text' => $new_content
							)
						)
					);
				}
			} else {
				$merged[]  = $message;
				$last_idx  = count( $merged ) - 1;
				$last_role = $message['role'];
			}
		}

		return $merged;
	}

	/**
	 * Prepares a single message based on its type.
	 *
	 * @param  mixed  $message  A message object (Message, Tool_Calls_Message, or Tool_Result_Message).
	 * @param  AI_Model  $model  The current model for compatibility checks.
	 *
	 * @return array|null The API-ready message structure or null to skip.
	 * @throws Exception If message type is unsupported.
	 * @since 1.0.9
	 */
	protected function prepare_individual_message( $message, AI_Model $model ) {
		if ( $message instanceof Tool_Calls_Message ) {
			return array(
				'role'    => 'assistant',
				'content' => $message->get_content_blocks(),
			);
		} elseif ( $message instanceof Tool_Result_Message ) {
			return array(
				'role'    => 'user',
				'content' => $message->get_content_blocks(),
			);
		} elseif ( $message instanceof Message ) {
			return ( new Message_Service() )->prepare_message( $message, $model );
		} else {
			throw new Exception( Error_Codes::NOT_SUPPORTED, __( 'Not supported message !', 'limb-chatbot' ) );
		}
	}

	/**
	 * Gets the max tokens value for the request.
	 *
	 * @return int The max tokens value.
	 * @since 1.0.9
	 */
	private function get_max_tokens(): int {
		// Check if global_utility exists and has get_max_completion_tokens method
		if ( property_exists( $this->utility, 'global_utility' )
		     && is_object( $this->utility->global_utility )
		     && method_exists( $this->utility->global_utility, 'get_max_completion_tokens' )
		) {
			$max_tokens = $this->utility->global_utility->get_max_completion_tokens();
			if ( ! empty( $max_tokens ) ) {
				return (int) $max_tokens;
			}
		}

		return self::DEFAULT_MAX_TOKENS;
	}

	/**
	 * Generates content and finalizes the response (with streaming support).
	 *
	 * @return Message|Tool_Calls_Message The generated message.
	 * @throws Exception If generation fails.
	 * @since 1.0.9
	 */
	public function finalize() {
		$model = $this->utility->global_utility->get_ai_model();
		if ( ! $model ) {
			throw new Exception( Error_Codes::AI_MODEL_NOT_SET, __( 'AI Model is not set.', 'limb-chatbot' ) );
		}
		$is_stream   = $this->utility->global_utility->get_stream();
		$messages    = $this->prepare_messages( $model );
		$http_client = $this->http_client_factory( $is_stream );

		$body = array(
			'model'      => $model->get_name(),
			'max_tokens' => $this->get_max_tokens(),
			'messages'   => $messages['messages'],
		);

		// Add system prompt if present
		if ( ! empty( $messages['system'] ) ) {
			$body['system'] = $messages['system'];
		}

		// Add tools if present
		$tools = $this->utility->get_tools();
		if ( ! empty( $tools ) ) {
			$body['tools'] = $tools;
		}

		// Set temperature
		$body['temperature'] = 0; // Make this dynamic

		// Enable streaming
		if ( $is_stream ) {
			$body['stream'] = true;
		}

		$args = array(
			'body'    => wp_json_encode( $body ),
			'headers' => $this->get_header(),
			'timeout' => $this->utility->get_timeout() ?? 120,
		);

		if ( $is_stream ) {
			$stream_parser         = new Stream_Parser();
			$args['stream_parser'] = array( $stream_parser, 'parser' );
		}

		$response = $http_client->post( self::API_BASE_URL . '/messages', $args );

		return ( new Chat_Completion_Response_Handler( $response, $http_client, $this,
			$stream_parser ?? null ) )->get_message();
	}

	/**
	 * Send a request to the AI endpoint.
	 *
	 * @param  Collection|array  $messages  Messages to send.
	 * @param  array  $params  Optional generation parameters.
	 *
	 * @return Chat_Completion_Response_Handler Response handler from HTTP client.
	 * @throws Exception If request fails.
	 * @since 1.0.9
	 */
	public function make_request( $messages, $stream = false, array $params = array() ) {
		$http_client = $this->http_client_factory($stream);
		$model       = null;

		if ( property_exists( $this->utility, 'ai_model' ) && ! empty( $this->utility->ai_model ) ) {
			$model = $this->utility->ai_model;
		}
		if ( empty( $model ) || ( ! $model instanceof AI_Model ) ) {
			$model = $this->utility->global_utility->get_ai_model();
		}

		// Prepare messages
		$prepared_messages = array();
		$system_content    = null;

		foreach ( $messages as $message ) {
			$prepared = $this->prepare_individual_message( $message, $model );
			if ( $prepared === null ) {
				continue;
			}

			if ( isset( $prepared['is_system'] ) && $prepared['is_system'] ) {
				$system_content = $prepared['content'];
				continue;
			}

			$prepared_messages[] = $prepared;
		}

		// Merge consecutive messages
		$prepared_messages = $this->merge_consecutive_messages( $prepared_messages );

		$body = array_merge(
			array(
				'model'      => $model->get_name(),
				'max_tokens' => $this->get_max_tokens(),
				'messages'   => $prepared_messages,
			),
			$params
		);

		if ( $stream ) {
			$body['stream']         = $stream;
		}

		if ( ! empty( $system_content ) ) {
			$body['system'] = $system_content;
		}

		$args = array(
			'body'    => wp_json_encode( $body ),
			'headers' => $this->get_header(),
			'timeout' => $this->utility->get_timeout() ?? 120,
		);

		if ( $stream ) {
			$stream_parser         = new Stream_Parser();
			$args['stream_parser'] = array( $stream_parser, 'parser' );
		}

		$response = $http_client->post( self::API_BASE_URL . '/messages', $args );

		return (new Chat_Completion_Response_Handler(
			$response,
			$http_client,
			$this,
			$stream_parser ?? null
		));
	}

	/**
	 * Generate Q&A entries using Claude.
	 *
	 * @param  Collection  $messages  Messages for Q&A generation.
	 * @param  int  $count  Number of entries to generate.
	 *
	 * @return Collection|null Generated Q&A dataset entries.
	 * @throws Exception If generation fails.
	 * @since 1.0.9
	 */
	public function generate_qa_entries( Collection $messages, int $count = 5 ): ?Collection {
		return $this->make_request( $messages, false, array( 'temperature' => 0.3 ) )->get_dataset_entries();
	}

	/**
	 * Check if vector search is applicable.
	 *
	 * @param  mixed  $messages  Collection of messages to process.
	 *
	 * @return Message The AI-generated response message.
	 * @throws Exception If request fails.
	 * @since 1.0.9
	 */
	public function is_vector_search_applicable( $messages ) {
		return $this->make_request( $messages, false, array( 'temperature' => 0.3 ) )->get_response_message();
	}

	/**
	 * Analyze a single turn of conversation.
	 *
	 * @param  Collection|array  $messages  Messages for analysis.
	 *
	 * @return array JSON-decoded response from AI endpoint.
	 * @throws Exception If request fails.
	 * @since 1.0.9
	 */
	public function analyze_turn( $messages ) {
		return $this->make_request( $messages, false, array( 'temperature' => 0.1 ) )->get_json();
	}

	/**
	 * Generate simple text response from AI.
	 *
	 * @param  Collection  $messages  Collection of Message objects for the AI.
	 * @param  bool  $stream  Whether to stream the response.
	 *
	 * @return Message The generated text response.
	 * @throws Exception If generation fails.
	 * @since 1.0.9
	 */
	public function generate_data_collect_question( $messages, $stream = false ) {
		return $this->make_request( $messages, $stream, array( 'temperature' => 0.7 ) )->get_message();
	}

	/**
	 * Generate a beautiful success message for action completion.
	 *
	 * @param  Collection|array  $messages  Messages for success message generation.
	 * @param  bool  $stream  Whether to stream the response.
	 *
	 * @return Message The generated success message text.
	 * @throws Exception If generation fails.
	 * @since 1.0.9
	 */
	public function generate_action_success_message( $messages, $stream = false ) {
		return $this->make_request( $messages, $stream, array( 'temperature' => 0.7 ) )->get_message();
	}

	/**
	 * Rerank messages based on relevance or quality.
	 *
	 * @param  Collection|array  $messages  Messages to rerank.
	 *
	 * @return Message The AI-generated message after reranking.
	 * @throws Exception If request fails.
	 * @since 1.0.9
	 */
	public function rerank( $messages ) {
		return $this->make_request( $messages, false, array( 'temperature' => 0 ) )->get_message();
	}

	/**
	 * Generate a JSON message response.
	 *
	 * @param  mixed  $messages  Messages to send.
	 *
	 * @return array JSON-decoded response.
	 * @throws Exception
	 * @since 1.0.9
	 */
	public function generate_json_message( $messages ) {
		return $this->make_request( $messages, false, array( 'temperature' => 0.6 ) )->get_json();
	}

	/**
	 * Generate a JSON response for reranking.
	 *
	 * @param  mixed  $messages  Messages to send.
	 *
	 * @return array JSON-decoded response.
	 * @throws Exception
	 * @since 1.0.9
	 */
	public function generate_rerank_json_message( $messages ) {
		return $this->make_request( $messages, false, array( 'temperature' => 0.2 ) )->get_json();
	}
}
