<?php
/**
 * Parses plain-text pseudo tool calls emitted by models into tool-call payloads.
 *
 * @package Niwa_AI
 */

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

class Niwa_AI_Agent_Pseudo_Tool_Parser {

	public static function extract_from_content( $content, $agent_context ) {
		$content = self::strip_code_fence( (string) $content );
		if ( '' === $content ) {
			return array();
		}

		if ( ! is_object( $agent_context ) || ! method_exists( $agent_context, 'get_tools' ) ) {
			return array();
		}

		$api_options = method_exists( $agent_context, 'get_api_options' ) ? $agent_context->get_api_options() : array();
		$tool_choice = isset( $api_options['tool_choice'] ) ? (string) $api_options['tool_choice'] : 'auto';
		if ( 'none' === $tool_choice ) {
			return array();
		}

		$valid_tools = self::available_function_tool_name_map( $agent_context );
		if ( empty( $valid_tools ) ) {
			return array();
		}

		$parse_content = $content;
		if ( ! self::content_starts_with_pseudo_tool_call( $parse_content, $valid_tools ) ) {
			$parse_content = self::extract_embedded_content( $content, $valid_tools );
			if ( '' === $parse_content ) {
				return array();
			}
		}

		$xml_calls = self::extract_xml_calls( $parse_content, $valid_tools );
		if ( ! empty( $xml_calls ) ) {
			return $xml_calls;
		}

		$json_call = self::extract_json_call( $parse_content, $valid_tools );
		if ( ! empty( $json_call ) ) {
			return array( $json_call );
		}

		return array();
	}

	public static function build_assistant_message( $tool_calls ) {
		$message = array(
			'role'       => 'assistant',
			'content'    => '',
			'tool_calls' => array(),
		);

		foreach ( (array) $tool_calls as $call ) {
			if ( ! is_array( $call ) ) {
				continue;
			}
			$name = isset( $call['name'] ) ? trim( (string) $call['name'] ) : '';
			if ( '' === $name ) {
				continue;
			}
			$id        = isset( $call['id'] ) ? trim( (string) $call['id'] ) : '';
			$arguments = isset( $call['arguments'] ) ? $call['arguments'] : '{}';
			if ( ! is_string( $arguments ) ) {
				$arguments = wp_json_encode( is_array( $arguments ) ? $arguments : array() );
			}
			if ( '' === $id ) {
				$id = self::generate_call_id();
			}
			if ( '' === $arguments ) {
				$arguments = '{}';
			}
			$message['tool_calls'][] = array(
				'id'       => $id,
				'type'     => 'function',
				'function' => array(
					'name'      => $name,
					'arguments' => $arguments,
				),
			);
		}

		return $message;
	}

	public static function should_retry_after_text( $stream_result, $agent_context ) {
		if ( ! is_array( $stream_result ) || empty( $stream_result['success'] ) ) {
			return false;
		}

		$content = isset( $stream_result['content'] ) ? trim( (string) $stream_result['content'] ) : '';
		if ( '' === $content ) {
			return false;
		}

		if ( ! is_object( $agent_context ) || ! method_exists( $agent_context, 'get_tools' ) ) {
			return false;
		}

		$tools = $agent_context->get_tools();
		if ( empty( $tools ) || ! is_array( $tools ) ) {
			return false;
		}

		$valid_tools = self::available_function_tool_name_map( $agent_context );
		if ( empty( $valid_tools ) || ! self::content_starts_with_pseudo_tool_call( $content, $valid_tools ) ) {
			return false;
		}

		$api_options = method_exists( $agent_context, 'get_api_options' ) ? $agent_context->get_api_options() : array();
		$tool_choice = isset( $api_options['tool_choice'] ) ? (string) $api_options['tool_choice'] : 'auto';
		return 'none' !== $tool_choice;
	}

	private static function extract_embedded_content( $content, $valid_tools ) {
		$content = self::strip_code_fence( (string) $content );
		if ( '' === $content || false === strpos( $content, '<' ) ) {
			return '';
		}

		if ( ! preg_match_all( '/<([a-z][a-z0-9_]*)\b/i', $content, $matches, PREG_SET_ORDER | PREG_OFFSET_CAPTURE ) ) {
			return '';
		}

		foreach ( $matches as $match ) {
			$tag = isset( $match[1][0] ) ? strtolower( trim( (string) $match[1][0] ) ) : '';
			if ( '' === $tag ) {
				continue;
			}

			if ( 'tool_call' === $tag || ( is_array( $valid_tools ) && ! empty( $valid_tools[ $tag ] ) ) ) {
				$offset = isset( $match[0][1] ) ? intval( $match[0][1] ) : 0;
				return ltrim( substr( $content, $offset ) );
			}
		}

		return '';
	}

	private static function extract_xml_calls( $content, $valid_tools ) {
		$content = self::strip_code_fence( (string) $content );
		if ( '' === $content || '<' !== substr( ltrim( $content ), 0, 1 ) ) {
			return array();
		}

		$calls     = array();
		$seen      = array();
		$remaining = ltrim( $content );
		for ( $i = 0; $i < 8; $i++ ) {
			$call = self::extract_leading_xml_call( $remaining, $valid_tools );
			if ( empty( $call ) ) {
				break;
			}

			$consumed = isset( $call['_consumed'] ) ? intval( $call['_consumed'] ) : 0;
			unset( $call['_consumed'] );

			$signature = ( isset( $call['name'] ) ? (string) $call['name'] : '' ) . "\n" . ( isset( $call['arguments'] ) ? (string) $call['arguments'] : '' );
			if ( empty( $seen[ $signature ] ) ) {
				$calls[]            = $call;
				$seen[ $signature ] = true;
			}

			if ( $consumed <= 0 ) {
				break;
			}

			$remaining = ltrim( substr( $remaining, $consumed ) );
			if ( '' === $remaining || '<' !== substr( $remaining, 0, 1 ) ) {
				break;
			}
		}

		return $calls;
	}

	private static function extract_leading_xml_call( $content, $valid_tools ) {
		$content = self::strip_code_fence( (string) $content );
		if ( '' === $content || '<' !== substr( ltrim( $content ), 0, 1 ) ) {
			return array();
		}

		$name     = '';
		$body     = '';
		$args     = array();
		$consumed = 0;
		if ( preg_match( '/^<tool_call\b([^>]*)>(.*?)<\/tool_call>/is', $content, $matches ) ) {
			$attrs    = self::parse_xml_attributes( $matches[1] );
			$name     = isset( $attrs['name'] ) ? (string) $attrs['name'] : '';
			$body     = isset( $matches[2] ) ? (string) $matches[2] : '';
			$consumed = strlen( $matches[0] );
		} elseif ( preg_match( '/^<([a-z0-9_]+)\b([^>]*)>(.*?)<\/\1>/is', $content, $matches ) ) {
			$name     = isset( $matches[1] ) ? (string) $matches[1] : '';
			$body     = isset( $matches[3] ) ? (string) $matches[3] : '';
			$args     = self::parse_xml_direct_arguments( isset( $matches[2] ) ? (string) $matches[2] : '' );
			$consumed = strlen( $matches[0] );
		} elseif ( preg_match( '/^<([a-z0-9_]+)\b(.*?)\/>/is', $content, $matches ) ) {
			$name     = isset( $matches[1] ) ? (string) $matches[1] : '';
			$args     = self::parse_xml_direct_arguments( isset( $matches[2] ) ? (string) $matches[2] : '' );
			$consumed = strlen( $matches[0] );
		}

		$name = trim( $name );
		if ( '' === $name || empty( $valid_tools[ $name ] ) ) {
			return array();
		}

		if ( preg_match_all( '/<param\b([^>]*)>(.*?)<\/param>/is', $body, $param_matches, PREG_SET_ORDER ) ) {
			foreach ( $param_matches as $param_match ) {
				$attrs      = self::parse_xml_attributes( isset( $param_match[1] ) ? (string) $param_match[1] : '' );
				$param_name = isset( $attrs['name'] ) ? trim( (string) $attrs['name'] ) : '';
				if ( '' === $param_name ) {
					continue;
				}
				$args[ $param_name ] = self::decode_argument_value( isset( $param_match[2] ) ? (string) $param_match[2] : '' );
			}
		}

		if ( empty( $args ) ) {
			$body_text = trim( html_entity_decode( $body, ENT_QUOTES | ENT_HTML5, 'UTF-8' ) );
			$body_text = self::strip_cdata_wrapper( $body_text );
			$json_len  = self::leading_balanced_json_object_length( $body_text );
			if ( $json_len > 0 ) {
				$decoded = json_decode( substr( $body_text, 0, $json_len ), true );
				if ( is_array( $decoded ) ) {
					$args = $decoded;
				}
			}
		}

		return array(
			'id'        => self::generate_call_id(),
			'name'      => $name,
			'arguments' => wp_json_encode( $args ),
			'_consumed' => $consumed,
		);
	}

	private static function extract_json_call( $content, $valid_tools ) {
		$content     = self::strip_code_fence( (string) $content );
		$json_length = self::leading_balanced_json_object_length( $content );
		if ( $json_length <= 0 ) {
			return array();
		}

		$decoded = json_decode( substr( $content, 0, $json_length ), true );
		if ( ! is_array( $decoded ) ) {
			return array();
		}

		$name = '';
		foreach ( array( 'tool_name', 'name', 'tool', 'function' ) as $key ) {
			if ( ! empty( $decoded[ $key ] ) && is_string( $decoded[ $key ] ) ) {
				$name = trim( (string) $decoded[ $key ] );
				break;
			}
		}

		if ( '' === $name || empty( $valid_tools[ $name ] ) ) {
			return array();
		}

		$args = array();
		foreach ( array( 'arguments', 'args', 'parameters', 'params' ) as $key ) {
			if ( array_key_exists( $key, $decoded ) ) {
				$args = self::normalize_arguments( $decoded[ $key ] );
				break;
			}
		}

		if ( empty( $args ) ) {
			$args = $decoded;
			foreach ( array( 'tool_name', 'name', 'tool', 'function', 'arguments', 'args', 'parameters', 'params' ) as $key ) {
				unset( $args[ $key ] );
			}
		}

		return array(
			'id'        => self::generate_call_id(),
			'name'      => $name,
			'arguments' => wp_json_encode( is_array( $args ) ? $args : array() ),
		);
	}

	private static function available_function_tool_name_map( $agent_context ) {
		if ( ! is_object( $agent_context ) || ! method_exists( $agent_context, 'get_tools' ) ) {
			return array();
		}

		$names = array();
		foreach ( (array) $agent_context->get_tools() as $tool ) {
			if ( ! is_array( $tool ) ) {
				continue;
			}

			$name = '';
			if ( isset( $tool['function']['name'] ) ) {
				$name = trim( (string) $tool['function']['name'] );
			} elseif ( isset( $tool['name'] ) && ( empty( $tool['type'] ) || 'function' === (string) $tool['type'] ) ) {
				$name = trim( (string) $tool['name'] );
			}

			if ( '' !== $name ) {
				$names[ $name ] = true;
			}
		}

		return $names;
	}

	private static function parse_xml_attributes( $attribute_text ) {
		$attrs = array();
		if ( ! preg_match_all( '/([a-zA-Z_:][-a-zA-Z0-9_:.]*)\s*=\s*(["\'])(.*?)\2/s', (string) $attribute_text, $matches, PREG_SET_ORDER ) ) {
			return $attrs;
		}
		foreach ( $matches as $match ) {
			$key = isset( $match[1] ) ? strtolower( (string) $match[1] ) : '';
			if ( '' === $key ) {
				continue;
			}
			$attrs[ $key ] = html_entity_decode( isset( $match[3] ) ? (string) $match[3] : '', ENT_QUOTES | ENT_HTML5, 'UTF-8' );
		}
		return $attrs;
	}

	private static function parse_xml_direct_arguments( $attribute_text ) {
		$attribute_text = (string) $attribute_text;
		$args           = array();
		if ( '' === trim( $attribute_text ) ) {
			return $args;
		}

		if ( ! preg_match_all( '/(?:^|\s)([a-zA-Z_:][-a-zA-Z0-9_:.]*)\s*=\s*(["\'])/s', $attribute_text, $matches, PREG_SET_ORDER | PREG_OFFSET_CAPTURE ) ) {
			return $args;
		}

		$count = count( $matches );
		for ( $i = 0; $i < $count; $i++ ) {
			$key = isset( $matches[ $i ][1][0] ) ? strtolower( (string) $matches[ $i ][1][0] ) : '';
			if ( '' === $key || 'name' === $key ) {
				continue;
			}

			$quote       = isset( $matches[ $i ][2][0] ) ? (string) $matches[ $i ][2][0] : '';
			$value_start = isset( $matches[ $i ][2][1] ) ? intval( $matches[ $i ][2][1] ) + strlen( $quote ) : 0;
			$value_end   = strlen( $attribute_text );
			if ( isset( $matches[ $i + 1 ][0][1] ) ) {
				$value_end = intval( $matches[ $i + 1 ][0][1] );
			}

			$value = trim( substr( $attribute_text, $value_start, max( 0, $value_end - $value_start ) ) );
			if ( '' !== $quote && substr( $value, -1 ) === $quote ) {
				$value = substr( $value, 0, -1 );
			}

			$args[ $key ] = self::decode_argument_value( $value );
		}

		return $args;
	}

	private static function normalize_arguments( $value ) {
		if ( is_array( $value ) ) {
			return $value;
		}
		if ( is_string( $value ) ) {
			$decoded = self::decode_argument_value( $value );
			return is_array( $decoded ) ? $decoded : array();
		}
		return array();
	}

	private static function decode_argument_value( $value ) {
		$value = trim( html_entity_decode( (string) $value, ENT_QUOTES | ENT_HTML5, 'UTF-8' ) );
		$value = self::strip_cdata_wrapper( $value );
		$value = self::strip_code_fence( $value );
		if ( '' === $value ) {
			return '';
		}

		$decoded = json_decode( $value, true );
		if ( JSON_ERROR_NONE === json_last_error() ) {
			return $decoded;
		}

		return $value;
	}

	private static function strip_cdata_wrapper( $value ) {
		$value = trim( (string) $value );
		if ( preg_match( '/^<!\[CDATA\[(.*)\]\]>$/s', $value, $matches ) ) {
			return trim( isset( $matches[1] ) ? (string) $matches[1] : '' );
		}
		return $value;
	}

	private static function strip_code_fence( $content ) {
		$content = trim( (string) $content );
		if ( preg_match( '/^```(?:json|xml)?\s*(.*?)\s*```$/is', $content, $matches ) ) {
			return trim( isset( $matches[1] ) ? (string) $matches[1] : '' );
		}
		return $content;
	}

	private static function generate_call_id() {
		if ( function_exists( 'wp_generate_uuid4' ) ) {
			return 'call_pseudo_' . str_replace( '-', '', wp_generate_uuid4() );
		}

		return 'call_pseudo_' . str_replace( '.', '', uniqid( '', true ) );
	}

	private static function content_starts_with_pseudo_tool_call( $content, $valid_tools = array() ) {
		$content = ltrim( (string) $content );
		if ( '' === $content ) {
			return false;
		}

		if ( '<' === substr( $content, 0, 1 ) ) {
			if ( preg_match( '/^<([a-z][a-z0-9_]*)\b/i', $content, $matches ) ) {
				$tag = isset( $matches[1] ) ? strtolower( trim( (string) $matches[1] ) ) : '';
				if ( 'tool_call' === $tag ) {
					return true;
				}
				return '' !== $tag && is_array( $valid_tools ) && ! empty( $valid_tools[ $tag ] );
			}

			return false;
		}

		if ( '{' !== substr( $content, 0, 1 ) ) {
			return false;
		}

		$json_length = self::leading_balanced_json_object_length( $content );
		if ( $json_length <= 0 ) {
			return false;
		}

		$decoded = json_decode( substr( $content, 0, $json_length ), true );
		if ( ! is_array( $decoded ) ) {
			return false;
		}

		foreach ( array( 'tool_name', 'name', 'arguments', 'args', 'cmd', 'command', 'argv', 'query' ) as $key ) {
			if ( array_key_exists( $key, $decoded ) ) {
				return true;
			}
		}

		return false;
	}

	private static function leading_balanced_json_object_length( $text ) {
		$text      = (string) $text;
		$length    = strlen( $text );
		$depth     = 0;
		$in_string = false;
		$escaped   = false;

		for ( $i = 0; $i < $length; $i++ ) {
			$char = $text[ $i ];

			if ( $in_string ) {
				if ( $escaped ) {
					$escaped = false;
				} elseif ( '\\' === $char ) {
					$escaped = true;
				} elseif ( '"' === $char ) {
					$in_string = false;
				}
				continue;
			}

			if ( '"' === $char ) {
				$in_string = true;
				continue;
			}

			if ( '{' === $char ) {
				$depth++;
				continue;
			}

			if ( '}' === $char ) {
				$depth--;
				if ( 0 === $depth ) {
					return $i + 1;
				}
			}
		}

		return 0;
	}
}
