<?php

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

use Limb_Chatbot\Includes\Services\Helper;
use Limb_Chatbot\Includes\Services\Stream_Event_Service;

/**
 * Class Stream_Parser
 *
 * Parses streamed chunks received from xAI's Chat Completions API.
 *
 * Handles reconstruction of incomplete JSON lines, identification of request metadata,
 * and manages stream state including streamed content output.
 *
 * xAI streaming uses Server-Sent Events (SSE) compatible with OpenAI format:
 * - data: JSON chunks with delta content
 * - [DONE] marker for stream end
 *
 * @package Limb_Chatbot\Includes\AI_Providers\Grok\Endpoints\Chat_Completion
 * @since 1.0.12
 */
class Stream_Parser {

	/**
	 * Collected JSON objects that were parsed completely from stream.
	 *
	 * @var array
	 * @since 1.0.12
	 */
	public array $complete_jsons = array();

	/**
	 * Flag to indicate if the first data chunk has started.
	 *
	 * @var bool
	 * @since 1.0.12
	 */
	public bool $data_chunk_started = false;

	/**
	 * Buffer for incomplete (half) JSON lines across multiple stream chunks.
	 *
	 * @var string
	 * @since 1.0.12
	 */
	public string $half_lines = '';

	/**
	 * Buffer for incomplete chunks (chunks split mid-line).
	 *
	 * @var string
	 * @since 1.0.12
	 */
	public string $chunk_buffer = '';

	/**
	 * Parsed error information from the stream, if any.
	 *
	 * @var array|null
	 * @since 1.0.12
	 */
	public ?array $error = null;

	/**
	 * Parses an individual stream chunk.
	 *
	 * Handles:
	 * - SSE data line detection
	 * - JSON line collection (including partial JSON across chunks)
	 * - Extraction of request-id for debugging
	 * - Stream content printing
	 * - Chunks split mid-line (accumulates in chunk_buffer)
	 *
	 * @param  string  $chunk  A single raw chunk received from the streaming API.
	 *
	 * @return void
	 * @since 1.0.12
	 */
	public function parser( $chunk ): void {
		// Add chunk to buffer - chunks may be split mid-line
		$this->chunk_buffer .= $chunk;

		// Split by newlines to process complete lines
		$lines = explode( "\n", $this->chunk_buffer );

		// Keep the last incomplete line in buffer for next chunk
		$this->chunk_buffer = array_pop( $lines );

		foreach ( $lines as $line ) {
			if ( empty( trim( $line ) ) ) {
				continue;
			}

			$line = trim( $line );

			// Handle request ID logging
			if ( str_contains( $line, 'request-id:' ) ) {
				$request_id = trim( explode( ':', $line, 2 )[1] ?? '' );
				if ( ! empty( $request_id ) ) {
					Helper::log( 'request-id: ' . $request_id );
				}
				continue;
			}

			// Handle data lines
			if ( str_starts_with( $line, 'data: ' ) ) {
				$line = substr( $line, 6 );
			}

			// Skip empty data or [DONE] marker
			if ( empty( trim( $line ) ) || $line === '[DONE]' ) {
				continue;
			}

			// Append to buffer (handles both raw JSON errors and SSE data)
			$this->half_lines .= $line;

			// Try to decode buffer if it looks like JSON
			if ( Helper::is_probable_json( $this->half_lines ) ) {
				$decoded = json_decode( $this->half_lines );

				if ( $decoded !== null ) {
					// Check for error response
					if ( isset( $decoded->error ) ) {
						$this->error = array(
							'type'    => $decoded->error->type ?? 'unknown',
							'message' => $decoded->error->message ?? 'Unknown error',
						);
						$this->half_lines = '';

						return;
					}

					$this->complete_jsons[] = $decoded;
					$this->print( $decoded );
					$this->half_lines = '';
				}
			}
		}
	}

	/**
	 * Prints streamed content to output.
	 *
	 * For xAI (OpenAI-compatible), we print delta content as it arrives.
	 *
	 * @param  object  $decoded  The decoded JSON object from the stream.
	 *
	 * @return void
	 * @since 1.0.12
	 */
	public function print( $decoded ): void {
		// Check if this is a chat completion chunk with choices
		if ( isset( $decoded->choices[0]->delta->content ) ) {
			$text = $decoded->choices[0]->delta->content;
			if ( ! empty( $text ) ) {
				if ( ! $this->data_chunk_started ) {
					$this->data_chunk_started = true;
				}
				// Send the chunk via the central service
				Stream_Event_Service::text( $text );
			}
		}
	}

	/**
	 * Gets the accumulated text content from all stream chunks.
	 *
	 * @return string The accumulated text content.
	 * @since 1.0.12
	 */
	public function get_text_content(): string {
		$content = '';

		foreach ( $this->complete_jsons as $chunk ) {
			if ( isset( $chunk->choices[0]->delta->content ) ) {
				$content .= $chunk->choices[0]->delta->content;
			}
		}

		return $content;
	}

	/**
	 * Gets tool calls from the stream chunks.
	 *
	 * @return array Array of tool calls.
	 * @since 1.0.12
	 */
	public function get_tool_calls(): array {
		$tool_calls = array();

		foreach ( $this->complete_jsons as $chunk ) {
			if ( isset( $chunk->choices[0]->delta->tool_calls ) ) {
				foreach ( $chunk->choices[0]->delta->tool_calls as $tool_call ) {
					$index = $tool_call->index ?? 0;

					if ( ! isset( $tool_calls[ $index ] ) ) {
						$tool_calls[ $index ] = array(
							'id'       => $tool_call->id ?? '',
							'type'     => $tool_call->type ?? 'function',
							'function' => array(
								'name'      => $tool_call->function->name ?? '',
								'arguments' => '',
							),
						);
					}

					if ( isset( $tool_call->function->arguments ) ) {
						$tool_calls[ $index ]['function']['arguments'] .= $tool_call->function->arguments;
					}
				}
			}
		}

		return array_values( $tool_calls );
	}

	/**
	 * Gets the stop reason from stream chunks.
	 *
	 * @return string|null The stop reason or null.
	 * @since 1.0.12
	 */
	public function get_stop_reason(): ?string {
		foreach ( array_reverse( $this->complete_jsons ) as $chunk ) {
			if ( isset( $chunk->choices[0]->finish_reason ) ) {
				return $chunk->choices[0]->finish_reason;
			}
		}

		return null;
	}

	/**
	 * Gets usage information from stream chunks.
	 *
	 * @return object|null Usage object or null.
	 * @since 1.0.12
	 */
	public function get_usage() {
		foreach ( array_reverse( $this->complete_jsons ) as $chunk ) {
			if ( isset( $chunk->usage ) ) {
				return $chunk->usage;
			}
		}

		return null;
	}
}
