<?php
/**
 * Chatbot tool-call orchestration.
 *
 * @package KimaAI\ChatBot\Tools
 */

declare(strict_types=1);

namespace KimaAI\ChatBot\Tools;

use KimaAI\Helper;

defined('ABSPATH') || exit;

/**
 * Runs generic tool-call extraction, execution, tracing, and continuation.
 */
final class ToolCallOrchestrator {

  /**
   * Create the orchestrator.
   *
   * @param ToolContinuationBuilder $continuation_builder Continuation payload builder.
   */
  public function __construct(private ToolContinuationBuilder $continuation_builder) {
  }

  /**
   * Register tool orchestration hooks.
   */
  public function registerHooks(): void {
    add_filter('kimaai_ai_chat_response', [$this, 'handleChatResponse'], 100, 4);
  }

  /**
   * Run the generic non-streaming tool-calling loop for supported providers.
   *
   * @param array $response Raw provider response.
   * @param string $provider Active provider.
   * @param array $payload Provider request payload.
   * @param mixed $ai_provider Provider client.
   * @return array
   */
  public function handleChatResponse(array $response, string $provider, array $payload, mixed $ai_provider): array {
    if (!in_array($provider, ['openai', 'openrouter', 'ollama'], true)) {
      return $response;
    }

    if (empty($payload['tools']) || !is_array($payload['tools'])) {
      return $response;
    }

    if (!is_object($ai_provider) || !method_exists($ai_provider, 'responses')) {
      return $response;
    }

    $endpoint = $ai_provider->responses();
    $current = $response;
    $trace = [];
    $history = is_array($payload['input'] ?? null) ? $payload['input'] : [];

    $max_rounds = 6;
    $round = 0;

    while ($round < $max_rounds) {
      ++$round;

      $calls = $this->extractFunctionCalls($current);
      if (!$calls) {
        break;
      }

      $execution = $this->executeFunctionCalls($calls, $payload, $round);
      $trace = array_merge($trace, $execution['trace']);

      $continuation = $this->continuation_builder->build($provider, $payload, $current, $execution['outputs'], $history);
      $history = $continuation['history'];
      $next = $continuation['payload'];

      if ($provider !== 'openrouter' && (string) ($next['previous_response_id'] ?? '') === '') {
        break;
      }

      $current = $provider === 'openrouter'
      ? $endpoint->create($next)
      : $endpoint->addContent((string) $next['previous_response_id'], $next);
    }

    if (!empty($trace)) {
      $current['_kimaai_tool_calls'] = $trace;
    }

    return $current;
  }

  /**
   * Extract function calls from direct output items and nested message content.
   *
   * @param array $response Raw provider response.
   * @return array<int, array{name: string, call_id: string, arguments: string, raw_call: array}>
   */
  public function extractFunctionCalls(array $response): array {
    $calls = [];

    foreach (($response['output'] ?? []) as $item) {
      if (!is_array($item)) {
        continue;
      }

      if (($item['type'] ?? null) === 'function_call') {
        $calls[] = [
          'name' => (string) ($item['name'] ?? ''),
          'call_id' => (string) (($item['call_id'] ?? '') ?: ($item['id'] ?? '')),
          'arguments' => (string) ($item['arguments'] ?? '{}'),
          'raw_call' => $item,
        ];
      }

      foreach (($item['content'] ?? []) as $part) {
        if (is_array($part) && ($part['type'] ?? null) === 'function_call') {
          $calls[] = [
            'name' => (string) ($part['name'] ?? ''),
            'call_id' => (string) (($part['call_id'] ?? '') ?: ($part['id'] ?? '')),
            'arguments' => (string) ($part['arguments'] ?? '{}'),
            'raw_call' => $part,
          ];
        }
      }
    }

    return $calls;
  }

  /**
   * Execute all function calls returned in one model response.
   *
   * @param array $calls Extracted function calls.
   * @param array $payload Chatbot or provider payload.
   * @param int $round Tool loop round.
   * @param callable|null $lifecycle Optional lifecycle callback.
   * @return array{outputs: array, trace: array}
   */
  public function executeFunctionCalls(array $calls, array $payload, int $round, ?callable $lifecycle = null): array {
    $outputs = [];
    $trace = [];

    foreach ($calls as $call) {
      $name = (string) ($call['name'] ?? '');
      $call_id = (string) ($call['call_id'] ?? '');
      $args_json = (string) ($call['arguments'] ?? '{}');
      $args = json_decode($args_json, true);
      $args = is_array($args) ? $args : [];
      $event_call_id = $call_id !== '' ? $call_id : 'missing_call_id';

      $this->emitLifecycle(
        $lifecycle,
        [
          'type' => 'tool',
          'status' => 'progress',
          'name' => $name !== '' ? $name : 'tool',
          'callId' => $event_call_id,
        ]
      );

      if ($call_id === '') {
        $output_item = [
          'type' => 'function_call_output',
          'call_id' => 'missing_call_id',
          'output' => wp_json_encode(['error' => 'Missing call_id for tool call'], JSON_UNESCAPED_UNICODE | JSON_UNESCAPED_SLASHES),
        ];

        $outputs[] = $output_item;
        $trace[] = $this->buildTraceItem(
          $round, 'error', $name, '', $args_json, $args, $call['raw_call'] ?? [], null, 'Missing call_id for tool call', $output_item
        );

        $this->emitLifecycle(
          $lifecycle,
          [
            'type' => 'tool',
            'status' => 'error',
            'name' => $name !== '' ? $name : 'tool',
            'callId' => $event_call_id,
          ]
        );

        continue;
      }

      try {
        $result = $this->executeTool($name, $args, $payload);
        $out_str = is_string($result)
        ? $result
        : wp_json_encode($result, JSON_UNESCAPED_UNICODE | JSON_UNESCAPED_SLASHES);

        $output_item = [
          'type' => 'function_call_output',
          'id' => 'fc_output_' . sanitize_key($call_id !== '' ? $call_id : uniqid('', true)),
          'call_id' => $call_id,
          'output' => (string) $out_str,
        ];

        $outputs[] = $output_item;
        $trace[] = $this->buildTraceItem($round, 'success', $name, $call_id, $args_json, $args, $call['raw_call'] ?? [], $result, null, $output_item);

        $this->emitLifecycle(
          $lifecycle,
          [
            'type' => 'tool',
            'status' => 'end',
            'name' => $name !== '' ? $name : 'tool',
            'callId' => $event_call_id,
          ]
        );
      } catch (\Throwable $exception) {
        $output_item = [
          'type' => 'function_call_output',
          'id' => 'fc_output_' . sanitize_key($call_id !== '' ? $call_id : uniqid('', true)),
          'call_id' => $call_id,
          'output' => wp_json_encode(['error' => $exception->getMessage()], JSON_UNESCAPED_UNICODE | JSON_UNESCAPED_SLASHES),
        ];

        $outputs[] = $output_item;
        $trace[] = $this->buildTraceItem(
          $round, 'error', $name, $call_id, $args_json, $args, $call['raw_call'] ?? [], null, $exception->getMessage(), $output_item
        );

        $this->emitLifecycle(
          $lifecycle,
          [
            'type' => 'tool',
            'status' => 'error',
            'name' => $name !== '' ? $name : 'tool',
            'callId' => $event_call_id,
          ]
        );
      }
    }

    return [
      'outputs' => $outputs,
      'trace' => $trace,
    ];
  }

  /**
   * Build the next streaming provider payload after tool outputs.
   *
   * @param string $provider Active provider.
   * @param array $provider_payload Current provider payload.
   * @param array $completed Completed provider response.
   * @param array $outputs Function call output items.
   * @return array{payload: array, history: array}
   */
  public function buildStreamingContinuation(string $provider, array $provider_payload, array $completed, array $outputs): array {
    return $this->continuation_builder->buildStreaming($provider, $provider_payload, $completed, $outputs);
  }

  /**
   * Execute one tool call through public extension filters.
   *
   * @param string $name Tool name.
   * @param array $args Decoded tool arguments.
   * @param array $payload Chatbot or provider payload.
   * @return mixed
   */
  private function executeTool(string $name, array $args, array $payload): mixed {
    $result = apply_filters('kimaai_chatbot_tool_call', null, $name, $args, $payload);

    if ($result !== null) {
      return $result;
    }

    $result = apply_filters('kimaai_stream_tool_call', null, $name, $args, $payload);

    if ($result !== null) {
      return $result;
    }

    return ['error' => 'Unknown tool: ' . $name];
  }

  /**
   * Build a JSON-safe trace entry for one tool call.
   *
   * @param int $round Tool loop round.
   * @param string $status Tool call status.
   * @param string $name Tool name.
   * @param string $call_id Provider call id.
   * @param string $args_json Raw JSON argument string.
   * @param array $args Decoded arguments.
   * @param array $raw_call Raw function call item.
   * @param mixed $result Successful tool result.
   * @param string|null $error Error message.
   * @param array $output_item Function call output item.
   * @return array
   */
  private function buildTraceItem(
    int $round,
    string $status,
    string $name,
    string $call_id,
    string $args_json,
    array $args,
    array $raw_call,
    mixed $result,
    ?string $error,
    array $output_item
  ): array {
    $item = [
      'round' => $round,
      'status' => $status,
      'name' => $name,
      'call_id' => $call_id,
      'arguments_json' => $args_json,
      'arguments' => $args,
      'raw_call' => Helper::normalizeToolTraceValue($raw_call),
      'output_item' => $output_item,
    ];

    if ($status === 'success') {
      $item['result'] = Helper::normalizeToolTraceValue($result);
    }
    else {
      $item['error'] = (string) $error;
    }

    return $item;
  }

  /**
   * Emit a streaming lifecycle event when a callback is available.
   *
   * @param callable|null $lifecycle Lifecycle callback.
   * @param array $event Event payload.
   */
  private function emitLifecycle(?callable $lifecycle, array $event): void {
    if ($lifecycle) {
      $lifecycle($event);
    }
  }
}
