<?php
/**
 * Streaming chatbot request runner.
 *
 * @package KimaAI\ChatBot
 */

declare(strict_types=1);

namespace KimaAI\ChatBot;

use KimaAI\Ai\AiService;
use KimaAI\Ai\Providers\Contracts\StreamingChatClientInterface;
use KimaAI\Ai\Providers\DTO\ChatRequest;
use KimaAI\Ai\Providers\DTO\StreamChunk;
use KimaAI\Ai\Providers\Factory;
use KimaAI\ChatBot\Tools\ToolCallOrchestrator;
use KimaAI\Logger;
use WP_REST_Request;
use WP_REST_Response;

defined('ABSPATH') || exit;

class ChatbotStreamingRequestRunner {

  /**
   * Create the runner.
   *
   * @param AiService $ai_service AI service.
   * @param ToolCallOrchestrator $tool_orchestrator Tool call orchestrator.
   * @param ChatbotProviderErrorNormalizer $error_normalizer Provider error normalizer.
   * @param ChatbotSseEmitter $emitter SSE emitter.
   */
  public function __construct(
    private AiService $ai_service,
    private ToolCallOrchestrator $tool_orchestrator,
    private ChatbotProviderErrorNormalizer $error_normalizer,
    private ChatbotSseEmitter $emitter
  ) {
  }

  /**
   * Return a copy of the runner using a different emitter.
   *
   * @param ChatbotSseEmitter $emitter Replacement SSE emitter.
   * @return self
   */
  public function withEmitter(ChatbotSseEmitter $emitter): self {
    return new self(
      $this->ai_service,
      $this->tool_orchestrator,
      $this->error_normalizer,
      $emitter
    );
  }

  /**
   * Handle a prepared streaming chatbot request.
   *
   * @param ChatbotPreparedRequest $prepared Prepared request.
   * @param WP_REST_Request $request REST request (unused; required by WP REST callback interface).
   * @return WP_REST_Response
   *
   * phpcs:disable -- $request is required by the WP REST callback interface.
   */
  public function handle(ChatbotPreparedRequest $prepared, WP_REST_Request $request): WP_REST_Response {
    // phpcs:enable
    $input_ai_provider = (string) ($prepared->ai_payload['aiProvider'] ?? 'default');
    $input_ai_model = (string) ($prepared->ai_payload['aiModel'] ?? '');
    $env_key = ($input_ai_provider !== 'default' ? $input_ai_provider : null) ?? 'chat';

    try {
      $env = $this->ai_service->env_resolver->resolve($env_key, $input_ai_model);
    } catch (\Throwable $exception) {
      $normalized = [
        'provider'    => 'openai',
        /* translators: %s: error message describing the failure. */
        'userMessage' => sprintf(__('Error resolving AI environment for streaming: %s', 'kimaai'), $exception->getMessage()),
        'message'     => $exception->getMessage(),
        'type'        => 'environment_resolution_error',
        'code'        => 'environment_resolution_error',
        'param'       => null,
        'status'      => 500,
        'retryable'   => false,
        'rawMessage'  => $exception->getMessage(),
      ];

      Logger::error($normalized['rawMessage']);
      return $this->makeErrorResponse($normalized);
    }

    if (!in_array(($env['provider'] ?? ''), ['openai', 'openrouter', 'ollama'], true)) {
      $this->emitUnsupportedStreamingProvider($env);
      return new WP_REST_Response(null, 200);
    }

    $this->streamChat($env, $prepared);

    return new WP_REST_Response(null, 200);
  }

  /**
   * Check whether the request qualifies for SSE streaming.
   *
   * @param ChatbotPreparedRequest $prepared Prepared request.
   * @param WP_REST_Request $request REST request.
   * @return bool
   */
  public function isEligible(ChatbotPreparedRequest $prepared, WP_REST_Request $request): bool {
    $data = $prepared->data;
    $stream_flag = !empty($data['stream']) ||
    (bool) $request->get_param('stream') ||
    ((string) $request->get_header('x-kimaai-stream') === '1');
    $accept = (string) $request->get_header('accept');
    $wants_sse = stripos($accept, 'text/event-stream') !== false;

    return !empty($prepared->chatbot_settings['streaming']) && $stream_flag && $wants_sse;
  }

  /**
   * Start the SSE stream and run the full chat loop.
   *
   * @param array $env Resolved AI environment.
   * @param ChatbotPreparedRequest $prepared Prepared request.
   * @return void
   */
  private function streamChat(array $env, ChatbotPreparedRequest $prepared): void {
    $this->emitter->start();
    $this->emitter->send(
      [
        'type'      => 'meta',
        'sessionId' => (string) ($prepared->payload['sessionId'] ?? ''),
      ],
      'meta'
    );

    $input_ai_provider = (string) ($prepared->ai_payload['aiProvider'] ?? 'default');
    $input_ai_model = (string) ($prepared->ai_payload['aiModel'] ?? '');

    $model = ($input_ai_provider !== 'default' && $input_ai_model !== '')
    ? $input_ai_model
    : (string) ($env['model'] ?? ($input_ai_model !== '' ? $input_ai_model : 'gpt-4o'));

    $provider_input = $prepared->ai_payload;
    $provider_input['aiModel'] = $model;
    $provider_payload_initial = $this->ai_service->resolveChatPayload(
      (string) $env['provider'],
      $model,
      $provider_input,
      false
    );

    $client = Factory::create(
      (string) $env['provider'],
      (string) ($env['apiKey'] ?? ''),
      null,
      $this->providerOptionsFromEnv($env)
    );

    if (!$client instanceof StreamingChatClientInterface) {
      $normalized = $this->error_normalizer->normalize(
        [
          'message' => __('The selected provider does not support streaming.', 'kimaai'),
          'type'    => 'unsupported_streaming',
          'code'    => 'unsupported_streaming',
          'status'  => 400,
        ],
        (string) ($env['provider'] ?? 'openai'),
        400
      );

      $this->emitter->error($normalized);
      $this->emitter->send(['type' => 'done'], 'done');
      $this->emitter->done();
      return;
    }

    $this->consumeStream($client, $env, $prepared, $provider_payload_initial);
  }

  /**
   * Emit an error for a provider that does not support streaming and terminate the stream.
   *
   * @param array $env Resolved AI environment.
   * @return void
   */
  private function emitUnsupportedStreamingProvider(array $env): void {
    $this->emitter->start();
    $this->emitNormalizedError(
      [
        'message' => __('The selected provider does not support streaming.', 'kimaai'),
        'type'    => 'unsupported_streaming',
        'code'    => 'unsupported_streaming',
        'status'  => 400,
      ],
      (string) ($env['provider'] ?? 'openai'),
      400
    );
  }

  /**
   * Run the multi-round streaming loop, handling tool calls between rounds.
   *
   * @param StreamingChatClientInterface $client Streaming chat client.
   * @param array $env Resolved AI environment.
   * @param ChatbotPreparedRequest $prepared Prepared request.
   * @param array $provider_payload_initial Initial provider payload.
   * @return void
   */
  private function consumeStream(
    StreamingChatClientInterface $client,
    array $env,
    ChatbotPreparedRequest $prepared,
    array $provider_payload_initial
  ): void {
    $final_response_obj = null;
    $tool_calls_trace = [];
    $provider_payload = $provider_payload_initial;

    $max_rounds = 6;
    $round = 0;

    while ($round < $max_rounds) {
      ++$round;
      $completed = $this->consumeRound($client, $env, $provider_payload);

      if (!is_array($completed)) {
        return;
      }

      $final_response_obj = $completed;
      $function_calls = $this->tool_orchestrator->extractFunctionCalls($completed);

      if (!$function_calls) {
        break;
      }

      $execution = $this->tool_orchestrator->executeFunctionCalls(
        $function_calls,
        $prepared->payload,
        $round,
        fn(array $event) => $this->emitter->send($event, 'tool')
      );

      $tool_calls_trace = array_merge($tool_calls_trace, $execution['trace']);
      $continuation = $this->tool_orchestrator->buildStreamingContinuation(
        (string) ($env['provider'] ?? 'openai'),
        $provider_payload,
        $completed,
        $execution['outputs']
      );
      $provider_payload = $continuation['payload'];
    }

    if (!is_array($final_response_obj)) {
      $this->emitNormalizedError(
        [
          'message' => __('No final response received.', 'kimaai'),
          'type'    => 'missing_final_response',
          'code'    => 'missing_final_response',
          'status'  => 502,
        ],
        (string) ($env['provider'] ?? 'openai'),
        502
      );
      return;
    }

    if (!empty($tool_calls_trace)) {
      $provider_payload_initial['tool_calls_trace'] = $tool_calls_trace;
    }

    $saved = $this->ai_service->saveChatResponse(
      (string) ($env['provider'] ?? 'openai'),
      $provider_payload_initial,
      $final_response_obj
    );

    $this->emitter->send(
      [
        'type'      => 'done',
        'sessionId' => (string) ($prepared->payload['sessionId'] ?? ''),
        'message'   => (string) ($saved['message'] ?? ''),
      ],
      'done'
    );
    $this->emitter->done();
  }

  /**
   * Consume one streaming round and return the completed response or null on error.
   *
   * @param StreamingChatClientInterface $client Streaming chat client.
   * @param array $env Resolved AI environment.
   * @param array $provider_payload Provider request payload.
   * @return array|null Completed response array, or null if the round ended with an error.
   */
  private function consumeRound(StreamingChatClientInterface $client, array $env, array $provider_payload): ?array {
    $completed = null;

    try {
      foreach ($client->streamChat(ChatRequest::fromArray($provider_payload)) as $chunk) {
        if (!$chunk instanceof StreamChunk) {
          continue;
        }

        switch ($chunk->type) {
          case StreamChunk::TYPE_TEXT_DELTA:
            $delta = (string) ($chunk->text ?? '');
            if ($delta !== '') {
              $this->emitter->send(
                  [
                    'type' => 'delta',
                    'text' => $delta,
                  ],
                  'delta'
                );
            }
              break;

          case StreamChunk::TYPE_TOOL:
            $this->forwardToolStreamChunk($chunk);
              break;

          case StreamChunk::TYPE_COMPLETED:
            $candidate = $chunk->data['response'] ?? ($chunk->raw['response'] ?? null);
            if (is_array($candidate)) {
              $completed = $candidate;
            }
              break;

          case StreamChunk::TYPE_ERROR:
            $error_source = $chunk->data['error'] ?? $chunk->raw ?? [
              'message' => __('Streaming failed.', 'kimaai'),
              'type'    => 'stream_error',
              'code'    => 'stream_error',
              'status'  => 502,
            ];

            if (is_array($error_source) && empty($error_source['message'])) {
              $error_source['message'] = __('Streaming failed.', 'kimaai');
            }

            $this->emitNormalizedError($error_source, (string) ($env['provider'] ?? 'openai'), 502);
              return null;

          case StreamChunk::TYPE_RAW:
          default:
              break;
        }
      }
    } catch (\Throwable $exception) {
      $this->emitNormalizedError($exception, (string) ($env['provider'] ?? 'openai'), 500);
      return null;
    }

    if (!is_array($completed)) {
      $this->emitNormalizedError(
        [
          'message' => __('Streaming ended without a completed response.', 'kimaai'),
          'type'    => 'stream_incomplete',
          'code'    => 'stream_incomplete',
          'status'  => 502,
        ],
        (string) ($env['provider'] ?? 'openai'),
        502
      );
      return null;
    }

    return $completed;
  }

  /**
   * Forward a tool stream chunk as a streaming event, skipping argument deltas.
   *
   * @param StreamChunk $chunk Tool stream chunk.
   * @return void
   */
  private function forwardToolStreamChunk(StreamChunk $chunk): void {
    $phase = (string) ($chunk->data['phase'] ?? '');

    if ($phase === '' || in_array($phase, ['arguments_delta', 'arguments_done'], true)) {
      return;
    }

    $status = match ($phase) {
      'added', 'in_progress' => 'progress',
      'searching' => 'progress',
      'completed' => 'end',
      'failed' => 'error',
      default => 'progress',
    };

    $event = [
      'type'   => 'tool',
      'status' => $status,
      'name'   => (string) ($chunk->data['name'] ?? 'tool'),
    ];

    $call_id = (string) (($chunk->data['call_id'] ?? '') ?: ($chunk->data['item_id'] ?? ''));
    if ($call_id !== '') {
      $event['callId'] = $call_id;
    }

    $this->emitter->send($event, 'tool');
  }

  /**
   * Normalize an error, emit it, and terminate the stream.
   *
   * @param mixed $error Raw error (Throwable, array, or string).
   * @param string $provider Active provider slug.
   * @param int $fallback_status HTTP fallback status code.
   * @return void
   */
  private function emitNormalizedError(mixed $error, string $provider, int $fallback_status): void {
    $normalized = $this->error_normalizer->normalize($error, $provider, $fallback_status);

    $this->emitter->error($normalized);
    $this->emitter->send(['type' => 'done'], 'done');
    $this->emitter->done();
  }

  /**
   * Build an error REST response from a normalized error payload.
   *
   * @param array $normalized Normalized error array.
   * @return WP_REST_Response
   */
  private function makeErrorResponse(array $normalized): WP_REST_Response {
    $status = (int) ($normalized['status'] ?? 500);

    return new WP_REST_Response(
      [
        'error'        => (string) ($normalized['userMessage'] ?? __('Something went wrong.', 'kimaai')),
        'errorDetails' => $normalized,
      ],
      $status
    );
  }

  /**
   * Build provider-specific options from an environment array.
   *
   * @param array $env Resolved AI environment.
   * @return array
   */
  private function providerOptionsFromEnv(array $env): array {
    $provider = strtolower(trim((string) ($env['provider'] ?? '')));
    $base_url = !empty($env['endpoint']) && is_string($env['endpoint'])
    ? trim($env['endpoint'])
    : '';

    if ($base_url === '') {
      return [];
    }

    if ($provider === 'ollama') {
      $parts = wp_parse_url($base_url);
      $path = isset($parts['path']) ? trim((string) $parts['path']) : '';

      if ($path === '' || $path === '/') {
        $base_url = rtrim($base_url, '/') . '/v1';
      }
    }

    return ['baseUrl' => rtrim($base_url, '/')];
  }
}
