<?php
/**
 * Handles partial Telegram turn-stream buffering and flush boundaries.
 *
 * @package Niwa_AI
 */

if ( ! defined( 'ABSPATH' ) ) {
    exit;
}

class Niwa_AI_Admin_Autonomy_Turn_Stream_Buffer {

    private $telegram;
    private $config;
    private $callbacks;

    public function __construct( $telegram, $config = array(), $callbacks = array() ) {
        $this->telegram  = $telegram;
        $this->config    = is_array( $config ) ? $config : array();
        $this->callbacks = is_array( $callbacks ) ? $callbacks : array();
    }

    public function maybe_flush( $state, $force = false ) {
        if ( ! is_array( $state ) || ! $this->telegram ) {
            return is_array( $state ) ? $state : array();
        }

        if ( ! empty( $state['assistant_stream_locked'] ) ) {
            if ( ! empty( $state['buffer'] ) ) {
                $this->log(
                    'turn_stream_flush_skipped_locked',
                    array(
                        'chat_id'    => isset( $state['chat_id'] ) ? (string) $state['chat_id'] : '',
                        'message_id' => isset( $state['message_id'] ) ? intval( $state['message_id'] ) : 0,
                        'buffer_len' => $this->text_length( (string) $state['buffer'] ),
                        'force'      => (bool) $force,
                    )
                );
                $state['buffer'] = '';
            }

            return $state;
        }

        $buffer     = isset( $state['buffer'] ) ? (string) $state['buffer'] : '';
        $buffer_len = $this->text_length( $buffer );
        if ( $this->starts_with_internal_protocol( $buffer ) ) {
            $this->log(
                $force ? 'turn_stream_internal_protocol_buffer_discarded' : 'turn_stream_internal_protocol_flush_blocked',
                array(
                    'chat_id'    => isset( $state['chat_id'] ) ? (string) $state['chat_id'] : '',
                    'message_id' => isset( $state['message_id'] ) ? intval( $state['message_id'] ) : 0,
                    'buffer_len' => $buffer_len,
                    'force'      => (bool) $force,
                )
            );
            if ( $force ) {
                $state['buffer'] = '';
            }

            return $state;
        }

        $should_flush = '' !== $buffer && $this->should_flush( $state, $force );
        if ( '' === $buffer || ! $should_flush ) {
            $started_at = isset( $state['started_at'] ) ? intval( $state['started_at'] ) : time();
            $elapsed    = max( 0, time() - $started_at );
            if ( ! $force && empty( $state['sent_any'] ) && $elapsed >= $this->config_int( 'initial_delay', 0 ) && $buffer_len >= 220 ) {
                $this->log(
                    'turn_stream_waiting_for_boundary',
                    array(
                        'chat_id'    => isset( $state['chat_id'] ) ? (string) $state['chat_id'] : '',
                        'message_id' => isset( $state['message_id'] ) ? intval( $state['message_id'] ) : 0,
                        'buffer_len' => $buffer_len,
                        'elapsed'    => $elapsed,
                        'force'      => (bool) $force,
                    )
                );
            }

            return $state;
        }

        $chat_id    = isset( $state['chat_id'] ) ? trim( (string) $state['chat_id'] ) : '';
        $message_id = isset( $state['message_id'] ) ? intval( $state['message_id'] ) : 0;
        if ( '' === $chat_id ) {
            return $state;
        }

        $payload    = $this->prepare_flush_payload( $state, $force );
        $flush_raw  = isset( $payload['flush_raw'] ) ? (string) $payload['flush_raw'] : '';
        $flush_text = isset( $payload['flush_text'] ) ? (string) $payload['flush_text'] : '';
        $remaining  = isset( $payload['remaining_buffer'] ) ? (string) $payload['remaining_buffer'] : '';
        if ( '' === $flush_text ) {
            $state['buffer'] = $remaining;
            return $state;
        }

        $this->log(
            'turn_stream_flush_attempt',
            array(
                'chat_id'       => $chat_id,
                'message_id'    => $message_id,
                'flush_len'     => $this->text_length( $flush_text ),
                'remaining_len' => $this->text_length( $remaining ),
                'force'         => (bool) $force,
                'sent_any'      => ! empty( $state['sent_any'] ),
                'sent_chars'    => intval( isset( $state['sent_chars'] ) ? $state['sent_chars'] : 0 ),
            )
        );

        $send_result = $this->telegram->send_message(
            $chat_id,
            $flush_text,
            ( empty( $state['sent_any'] ) && ! empty( $state['quote_current_message'] ) ) ? $message_id : 0
        );
        if ( empty( $send_result['success'] ) ) {
            $state['streaming_disabled'] = true;
            $this->log(
                'turn_stream_flush_failed',
                array(
                    'chat_id'     => $chat_id,
                    'message_id'  => $message_id,
                    'flush_len'   => $this->text_length( $flush_text ),
                    'send_result' => $send_result,
                )
            );

            return $state;
        }

        $state['sent_any']        = true;
        $state['sent_text']       = isset( $state['sent_text'] ) ? (string) $state['sent_text'] . $flush_raw : $flush_raw;
        $state['sent_chars']      = intval( isset( $state['sent_chars'] ) ? $state['sent_chars'] : 0 ) + $this->text_length( $flush_raw );
        $state['buffer']          = $remaining;
        $state['last_partial_at'] = time();
        $this->log(
            'turn_stream_flush_sent',
            array(
                'chat_id'       => $chat_id,
                'message_id'    => $message_id,
                'flush_len'     => $this->text_length( $flush_text ),
                'remaining_len' => $this->text_length( $remaining ),
                'sent_chars'    => intval( $state['sent_chars'] ),
                'send_result'   => $send_result,
            )
        );

        return $state;
    }

    public function starts_with_internal_protocol( $buffer ) {
        $buffer = ltrim( (string) $buffer );
        if ( '' === $buffer || '<' !== substr( $buffer, 0, 1 ) ) {
            return false;
        }

        $tag = $this->extract_leading_xml_tag_name( $buffer );
        if ( '' === $tag ) {
            return false;
        }
        if ( in_array( $tag, array( 'tool_call', 'arguments' ), true ) ) {
            return true;
        }
        if ( strlen( $tag ) >= 7 && 0 === substr_compare( $tag, '_result', -7 ) ) {
            return true;
        }

        $open_tag_end = strpos( $buffer, '>' );
        if ( false === $open_tag_end ) {
            return $this->is_structured_protocol_tag( $tag );
        }

        $is_structured_tag = $this->is_structured_protocol_tag( $tag );
        $open_tag          = substr( $buffer, 0, $open_tag_end + 1 );
        if ( $is_structured_tag && preg_match( '/\b(?:tool_call_id|status)\s*=/i', $open_tag ) ) {
            return true;
        }

        $body = ltrim( substr( $buffer, $open_tag_end + 1 ) );
        return $is_structured_tag && ( '' === $body || '{' === substr( $body, 0, 1 ) );
    }

    public function extract_leading_xml_tag_name( $buffer ) {
        if ( preg_match( '/^<([a-z][a-z0-9_]*)\b/i', ltrim( (string) $buffer ), $matches ) ) {
            return strtolower( (string) $matches[1] );
        }

        return '';
    }

    public function is_structured_protocol_tag( $tag ) {
        $tag = (string) $tag;
        return '' !== $tag && false !== strpos( $tag, '_' );
    }

    public function prepare_flush_payload( $state, $force = false ) {
        $state  = is_array( $state ) ? $state : array();
        $buffer = isset( $state['buffer'] ) ? (string) $state['buffer'] : '';
        $length = $this->text_length( $buffer );
        if ( $length <= 0 ) {
            return array(
                'flush_raw'        => '',
                'flush_text'       => '',
                'remaining_buffer' => '',
            );
        }

        $flush_raw = $buffer;
        $remaining = '';
        if ( ! $force ) {
            $min_split_length = ! empty( $state['sent_any'] )
                ? $this->config_int( 'followup_delay_min_chars', 0 )
                : 220;
            $split_length = $this->find_split_length( $buffer, $min_split_length );
            if ( $split_length > 0 && $split_length < $length ) {
                $flush_raw = $this->text_substr( $buffer, 0, $split_length );
                $remaining = $this->text_substr( $buffer, $split_length );
            }
        }

        return array(
            'flush_raw'        => $flush_raw,
            'flush_text'       => trim( $flush_raw ),
            'remaining_buffer' => $remaining,
        );
    }

    public function find_split_length( $buffer, $min_chars = 0 ) {
        $buffer = (string) $buffer;
        $length = $this->text_length( $buffer );
        if ( $length <= max( 0, intval( $min_chars ) ) ) {
            return 0;
        }

        $chars = preg_split( '//u', $buffer, -1, PREG_SPLIT_NO_EMPTY );
        if ( ! is_array( $chars ) || empty( $chars ) ) {
            return 0;
        }

        $min_index = max( 0, min( count( $chars ) - 1, intval( $min_chars ) - 1 ) );
        for ( $i = count( $chars ) - 1; $i >= $min_index; $i-- ) {
            if ( $this->is_hard_boundary_char( $chars[ $i ] ) ) {
                return $this->advance_split_length( $chars, $i + 1 );
            }
        }
        for ( $i = count( $chars ) - 1; $i >= $min_index; $i-- ) {
            if ( preg_match( '/\s/u', $chars[ $i ] ) ) {
                return $this->advance_split_length( $chars, $i + 1 );
            }
        }

        return 0;
    }

    public function advance_split_length( $chars, $length ) {
        $count  = is_array( $chars ) ? count( $chars ) : 0;
        $length = max( 0, intval( $length ) );
        while ( $length < $count ) {
            $char = (string) $chars[ $length ];
            if ( ! preg_match( '/[\s"\x{201D}\x{2019})\]]/u', $char ) ) {
                break;
            }
            $length++;
        }

        return $length;
    }

    public function is_hard_boundary_char( $char ) {
        return in_array( (string) $char, array( '.', '!', '?', ':', ';', "\n" ), true );
    }

    public function should_flush( $state, $force = false ) {
        if ( ! is_array( $state ) ) {
            return false;
        }

        $buffer = isset( $state['buffer'] ) ? (string) $state['buffer'] : '';
        $length = $this->text_length( $buffer );
        if ( $length <= 0 ) {
            return false;
        }
        if ( $force ) {
            return true;
        }

        $has_boundary = $this->has_flush_boundary( $buffer, $length );
        $started_at   = isset( $state['started_at'] ) ? intval( $state['started_at'] ) : time();
        $elapsed      = max( 0, time() - $started_at );
        if ( empty( $state['sent_any'] ) && $elapsed >= $this->config_int( 'initial_delay', 0 ) && $length >= 220 ) {
            return $has_boundary || $length >= $this->config_int( 'initial_hard_chars', 0 );
        }
        if ( $length >= $this->config_int( 'max_buffer_chars', PHP_INT_MAX ) ) {
            return true;
        }
        if ( ! empty( $state['sent_any'] ) ) {
            $last_partial_at    = isset( $state['last_partial_at'] ) ? intval( $state['last_partial_at'] ) : 0;
            $since_last_partial = $last_partial_at > 0 ? max( 0, time() - $last_partial_at ) : $elapsed;
            if ( $has_boundary && $length >= $this->config_int( 'followup_min_chars', 0 ) ) {
                return true;
            }
            if ( $since_last_partial >= $this->config_int( 'followup_delay', 0 ) && $length >= $this->config_int( 'followup_delay_min_chars', 0 ) ) {
                return true;
            }
        }
        if ( $length < $this->config_int( 'min_chars', 0 ) ) {
            return false;
        }

        return $has_boundary;
    }

    public function has_flush_boundary( $buffer, $length ) {
        $tail = trim( $this->text_substr( (string) $buffer, max( 0, intval( $length ) - 160 ) ) );
        return (bool) preg_match( '/(?:[.!?:]\s*|\n{2,})$/u', $tail );
    }

    private function config_int( $key, $default ) {
        return isset( $this->config[ $key ] ) ? intval( $this->config[ $key ] ) : intval( $default );
    }

    private function log( $event, $context = array() ) {
        $this->call( 'log', array( $event, is_array( $context ) ? $context : array() ), null );
    }

    private function text_length( $text ) {
        return intval( $this->call( 'text_length', array( $text ), strlen( (string) $text ) ) );
    }

    private function text_substr( $text, $start, $length = null ) {
        return (string) $this->call( 'text_substr', array( $text, $start, $length ), null === $length ? substr( (string) $text, intval( $start ) ) : substr( (string) $text, intval( $start ), intval( $length ) ) );
    }

    private function call( $name, $args = array(), $default = null ) {
        return Niwa_AI_Admin_Autonomy_Callbacks::call( $this->callbacks, $name, $args, $default );
    }
}
