<?php

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

use Limb_Chatbot\Includes\AI_Providers\Grok\Endpoints\Chat_Completion\Handlers\Chat_Completion_Response_Handler;
use Limb_Chatbot\Includes\AI_Providers\Grok\Endpoints\Grok_Endpoint;
use Limb_Chatbot\Includes\AI_Providers\Grok\Services\Message_Service;
use Limb_Chatbot\Includes\AI_Providers\Grok\Utilities\Chatbot_Utility;
use Limb_Chatbot\Includes\AI_Providers\Grok\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 xAI Chat Completions API interactions for assistant responses.
 *
 * @package Limb_Chatbot\Includes\AI_Providers\Grok\Endpoints\Chat_Completion
 * @since 1.0.12
 */
class Chat_Completion_Endpoint extends Grok_Endpoint {

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

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

	/**
	 * Sends a chat completions request to the xAI 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 xAI.
	 * @throws Exception If messages are empty or invalid.
	 * @since 1.0.12
	 */
	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,
		);

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

		// Set temperature
		$body['temperature'] = 0;

		// 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 . '/chat/completions', $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.
	 *
	 * @param  AI_Model  $model  The model for message formatting.
	 *
	 * @return array Array of formatted messages.
	 * @throws Exception If the message list is empty.
	 * @since 1.0.12
	 */
	private function prepare_messages( $model ) {
		$messages = array();

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

			$messages[] = $prepared;
		}

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

		return $messages;
	}

	/**
	 * 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.12
	 */
	protected function prepare_individual_message( $message, AI_Model $model ) {
		if ( $message instanceof Tool_Calls_Message ) {
			return array(
				'role'       => 'assistant',
				'tool_calls' => $message->get_tool_calls_array(),
			);
		} elseif ( $message instanceof Tool_Result_Message ) {
			return array(
				'role'         => 'tool',
				'tool_call_id' => $message->get_tool_call_id(),
				'content'      => $message->get_content_string(),
			);
		} 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.12
	 */
	private function get_max_tokens(): int {
		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.12
	 */
	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,
		);

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

		// Set temperature
		$body['temperature'] = 0;

		// 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 . '/chat/completions', $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  bool  $stream  Whether to stream the response.
	 * @param  array  $params  Optional generation parameters.
	 *
	 * @return Chat_Completion_Response_Handler Response handler from HTTP client.
	 * @throws Exception If request fails.
	 * @since 1.0.12
	 */
	/**
	 * Analyze a single turn of conversation.
	 *
	 * @param  Collection|array  $messages  Messages for analysis.
	 *
	 * @return array JSON-decoded response from AI endpoint.
	 * @throws Exception
	 * @since 1.0.12
	 */
	public function analyze_turn( $messages ) {
		return $this->make_request( $messages, false, array( 'temperature' => 0.1 ) )->get_json();
	}

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

	/**
	 * Generate a question for data collection.
	 *
	 * @param  mixed  $messages  Collection of Message objects for the AI.
	 * @param  bool  $stream  Whether to stream the response.
	 *
	 * @return Message The generated text response.
	 * @throws Exception
	 * @since 1.0.12
	 */
	public function generate_data_collect_question( $messages, $stream = false ) {
		return $this->make_request( $messages, $stream, array( 'temperature' => 0.7 ) )->get_message();
	}

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

	/**
	 * Generate a JSON message response.
	 *
	 * @param  mixed  $messages  Messages to send.
	 *
	 * @return array JSON-decoded response.
	 * @throws Exception
	 * @since 1.0.12
	 */
	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.12
	 */
	public function generate_rerank_json_message( $messages ) {
		return $this->make_request( $messages, false, array( 'temperature' => 0.2 ) )->get_json();
	}

	/**
	 * Rerank messages based on relevance.
	 *
	 * @param  mixed  $messages  Messages to rerank.
	 *
	 * @return Message The AI-generated reranked message.
	 * @throws Exception
	 * @since 1.0.12
	 */
	public function rerank( $messages ) {
		return $this->make_request( $messages, false, array( 'temperature' => 0 ) )->get_message();
	}

	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();

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

			$prepared_messages[] = $prepared;
		}

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

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

		$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 . '/chat/completions', $args );

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