<?php
/**
 * AiService class.
 *
 * @class AiService
 * @package KimaAI\Ai
 * @since 1.1.0
 */

declare(strict_types=1);

namespace KimaAI\Ai;

defined('ABSPATH') || exit;

use KimaAI\Admin\Settings;
use KimaAI\Ai\AiModels;
use KimaAI\Ai\Providers\Factory;
use KimaAI\Ai\Providers\Contracts\ChatClientInterface;
use KimaAI\Ai\Providers\Contracts\EmbeddingClientInterface;
use KimaAI\Ai\Providers\Contracts\ImageClientInterface;
use KimaAI\Ai\Providers\Contracts\ModelClientInterface;
use KimaAI\Ai\Providers\DTO\ChatRequest;
use KimaAI\Ai\Providers\DTO\EmbeddingRequest;
use KimaAI\Ai\Providers\DTO\ImageRequest;
use KimaAI\Helper;
use KimaAI\Logger;
use KimaAI\Models\ChatModel;
use KimaAI\Models\ImageModel;
use RuntimeException;

class AiService {

  /**
   * Resolved provider slug for the current request.
   *
   * @var string|null
   */
  private ?string $ai_provider = null;
  /**
   * Resolved environment identifier for the current request.
   *
   * @var string|null
   */
  private ?string $ai_env = null;
  /**
   * Active model identifier for the current request.
   *
   * @var string|null
   */
  private ?string $ai_model = null;
  /**
   * Active generation type for the current request.
   *
   * @var string|null
   */
  private ?string $generation_type = null;

  /**
   * Create the AI service facade.
   */
  public function __construct(
    public EnvResolver $env_resolver,
    public Settings $settings,
    public ImageModel $image_model,
    public ChatModel $chat_model,
  ) {
  }


  /**
   * Return a cloned AiService with a specific env/model/generation_type
   * for this job (e.g. Article Writer, Playground, etc.).
   *
   * $env_or_alias can be:
   *  - "default", "chat", "image", "embed"
   *  - "openai-XXXX", "replicate-XXXX" (env slug)
   *  - or raw env id "6872453c84..."
   */
  public function forEnv(
    ?string $env_or_alias = null,
    ?string $model = null,
    ?string $generation_type = null
  ): self {
    $clone = clone $this;

    if (!empty($env_or_alias)) {
      $clone->ai_env = $env_or_alias;
    }

    if (!empty($model)) {
      $clone->ai_model = $model;
    }

    if (!empty($generation_type)) {
      $clone->generation_type = $generation_type;
    }

    return $clone;
  }

  /**
   * Fire a one-off request to an AI service (no session / history kept).
   *
   * The $ai_env argument may be:
   *   – a raw env ID   (e.g. "6872453c84041")
   *   – a "provider-id" string (e.g. "openai-6872453c84041")
   *   – one of the logical aliases: "default", "chat", "image", "embed"
   *
   * @param array $input The input payload expected by the model.
   * @param string|null $ai_env Environment alias / ID (see above).
   *
   * @throws RuntimeException On any missing config or unknown provider.
   */
  public function singleCall(array $input, ?string $ai_env = null): string {
    $requested_env = $this->ai_env ?? $ai_env ?? 'default';
    $env = $this->env_resolver->resolve($requested_env);

    $provider = (string) ($env['provider'] ?? '');
    $api_key = (string) ($env['apiKey'] ?? '');

    $this->ai_model = $this->ai_model ?? (string) ($env['model'] ?? '');
    $model = (string) $this->ai_model;
    $this->ai_provider = $provider;
    $this->ai_env = $env['env_id'] ?? $requested_env;
    $this->generation_type = 'chat';

    // Fallback if env returned an empty model.
    if ($model === '') {
      $model = $provider === 'replicate' ? 'meta/meta-llama-3-70b-instruct' : 'gpt-4o';
      $this->ai_model = $model;
    }

    // Let callers override model in the payload; otherwise use env model.
    if (empty($input['model'])) {
      $input['model'] = $model;
    }

    $effective_model = (string) ($input['model'] ?: $model);
    $input = $this->applyPresetParams($env, $provider, $effective_model, $input);

    $ai_provider = Factory::create($provider, $api_key, null, $this->providerOptionsFromEnv($env));
    // New provider-layer-first path.
    if ($ai_provider instanceof ChatClientInterface) {
      return $ai_provider->createChat(ChatRequest::fromArray($input))->text;
    }
    $error_message = sprintf('Provider "%s" does not support chat.', esc_html($provider));
    Logger::error($error_message);
    throw new RuntimeException($error_message); // phpcs:ignore
  }

  /**
   * Creates an AI provider based on the input and generation type.
   *
   * @param array $input The input data.
   * @param string $generation_type The type of generation (e.g., 'chat', 'image', 'embed').
   * @param bool $is_playground Whether the request is from the playground.
   *
   * @return array The response data.
   *
   * @throws RuntimeException If the AI provider is unknown or the generation type is not supported.
   */
  public function run(array $input, string $generation_type, bool $is_playground = false): array {
    $this->generation_type = $generation_type;
    $input_ai_provider = $input['aiProvider'] ?? null;
    $input_ai_model = $input['aiModel'] ?? $input['model'] ?? null;

    // Use explicit env if provider ≠ "default"; otherwise fall back to the logical alias (chat/image/embed).
    $env_key = ($input_ai_provider !== 'default' ? $input_ai_provider : null) ?? $generation_type;

    $env = $this->env_resolver->resolve($env_key, $input_ai_model);
    $this->ai_provider = $env['provider'];
    $this->ai_env = $env['env_id'] ?? $env_key;
    $this->ai_model = ($input_ai_provider !== 'default' && $input_ai_model)
    ? $input_ai_model
    : ($this->ai_model ?? $env['model']);

    $ai_provider = Factory::create($this->ai_provider, $env['apiKey'], null, $this->providerOptionsFromEnv($env));

    if ($generation_type === 'chat') {
      $input = $this->applyPresetParams($env, $this->ai_provider, $this->ai_model, $input);
    }

    // Add the ai_provider object to the input array.
    $input['aiProviderObject'] = $ai_provider;

    // Mark that the default image provider is being used (for special handling).
    if ($generation_type === 'image' && ($input['aiProvider'] ?? 'default') === 'default') {
      $input['is_default'] = true;
    }

    // Build provider-specific payload.
    if ($generation_type !== 'embed') {
      $payload = $this->resolveChatPayload($this->ai_provider, $this->ai_model, $input, $is_playground);
    }
    else {
      // Embeddings use a simpler payload.
      $input['model'] = $this->ai_model;
      $payload = $input;
    }

    switch ($generation_type) {
      case 'chat':
        $chat_response = $ai_provider instanceof ChatClientInterface
        ? $ai_provider->createChat(ChatRequest::fromArray($payload))
        : null;
        $response = $chat_response ? $chat_response->raw : $ai_provider->chat($payload);

        /**
         * Filter: mutate / post-process the raw provider chat response.
         * Can be used to implement provider-specific function-calling loops, logging, etc.
         *
         * @param array $response Raw provider response.
         * @param string $provider 'openai', 'replicate', ...
         * @param array $payload Provider payload.
         * @param mixed $ai_provider The provider instance (OpenAi, Replicate, ...).
         */
        $response = apply_filters('kimaai_ai_chat_response', $response, $this->ai_provider, $payload, $ai_provider);

        $save_response = $this->saveChatResponse($this->ai_provider, $payload, $response);
          return $save_response;

      case 'image':
        $image_response = $ai_provider instanceof ImageClientInterface
        ? $ai_provider->createImage(ImageRequest::fromArray($payload))
        : null;
        $response = $image_response ? $image_response->raw : $ai_provider->image($payload);
        $save_response = $this->saveImageResponse($this->ai_provider, $this->ai_model, $payload, $response);
          return $save_response;

      case 'embed':
        $embedding_response = $ai_provider instanceof EmbeddingClientInterface
        ? $ai_provider->createEmbedding(EmbeddingRequest::fromArray($payload))
        : null;
        $response = $embedding_response ? $embedding_response->raw : $ai_provider->embedding($payload);
          return $response;

      default:
        $error_message = 'Unknown generation type: ' . esc_html($generation_type) . '.';
        Logger::error($error_message);
        throw new RuntimeException($error_message); // phpcs:ignore
    }
  }

  /**
   * Resolves the payload/query for chat-like requests based on the input
   * and the active provider.
   *
   * @param string $provider 'openai' or 'replicate'.
   * @param string $model Model name.
   * @param array $input The input data.
   * @param bool $is_playground Whether the request is from the playground.
   *
   * @return array The resolved payload.
   */
  public function resolveChatPayload(string $provider, string $model, array $input, bool $is_playground = false): array {
    $session_id = $input['sessionId'] ?? Helper::ensureSessionId($is_playground);
    $chat_history = $this->chat_model->getSessionHistory($session_id);
    $kb_context = $input['kbContext'] ?? null;

    $messages = Helper::buildMessageTree(
      $chat_history,
      $input['prompt'],
      $is_playground,
      $input['chatbot_settings'] ?? null,
      $input['postId'] ?? 0,
      $kb_context,
    );

    $rename_keys_for_replicate = [
      'n' => 'num_outputs',
      'responseFormat' => 'output_format',
    ];

    $payload = [];

    // Provider-specific key remapping.
    if ($provider === 'replicate') {
      foreach ($input as $key => $value) {
        if (isset($rename_keys_for_replicate[$key])) {
          $payload[$rename_keys_for_replicate[$key]] = $value;
        }
      }
    }

    // Common: camelCase → snake_case passthrough for known keys.
    foreach ($input as $key => $value) {
      $snake = strtolower(preg_replace('/([a-z])([A-Z])/', '$1_$2', $key));
      $payload[$snake] = $value;
    }

    $payload['is_playground'] = $is_playground;
    $payload['session_id'] = $session_id;
    $payload['model'] = $model;
    $payload['ai_provider'] = $provider;

    // For replicate compatibility.
    $payload['user_message'] = $input['prompt'] ?? '';

    if ($provider === 'replicate') {
      if (isset($input['is_default'])) {
        $payload['output_format'] = 'png';
      }

      $payload = array_replace(
        $payload,
        $this->formatMessagesForReplicate($model, $messages)
      );
    }
    elseif (in_array($provider, ['openai', 'openrouter', 'ollama'], true)) {
      $payload['input'] = $messages;
    }

    return $payload;
  }

  /**
   * Formats messages for the Replicate API.
   *
   * @param string $model The active model.
   * @param array $messages The messages to format.
   *
   * @return array The formatted messages.
   */
  public function formatMessagesForReplicate(string $model, array $messages): array {
    switch ($model) {
      // 1) Flux-Schnell (history + prompt split).
      case 'black-forest-labs/flux-schnell':
        $history = [];
        foreach ($messages as $message) {
          if ($message['role'] === 'system') {
            continue; // this model ignores the system prompt.
          }
          $history[] = $message['content'];
        }
        $prompt = array_pop($history); // last user turn.
          return compact('history', 'prompt');

      // 2) Models that already accept ChatML directly.
      case (str_contains($model, 'Hermes')):
          return ['messages' => $messages];

      // 3) Fallback: flatten everything into a single prompt.
      default:
        $prompt = implode(
            "\n",
            array_map(
              static fn($message) => strtoupper($message['role']) . ': ' . $message['content'],
              $messages
            )
          ) . "\nASSISTANT:";
          return ['prompt' => $prompt];
    }
  }

  /**
   * Save the chat response to the database.
   *
   * @param string $provider 'openai' or 'replicate'.
   * @param array $payload The payload data.
   * @param array $response The response data.
   *
   * @return array The saved response data.
   */
  public function saveChatResponse(string $provider, array $payload, array $response): array {
    $ai_response = $this->parseResponse($provider, $response);

    $ip_address = isset($_SERVER['REMOTE_ADDR'])
    ? filter_var(wp_unslash($_SERVER['REMOTE_ADDR']), FILTER_VALIDATE_IP)
    : '';

    // Save env_id inside the payload for auditability.
    $payload['env_id'] = $this->ai_env;

    // Save tool/function call trace inside the final payload row if present.
    if (!empty($payload['tool_calls_trace']) && is_array($payload['tool_calls_trace'])) {
      $tool_calls_trace = $payload['tool_calls_trace'];
    }
    else {
      $tool_calls_trace = $this->extractToolCallsTraceFromResponse($provider, $response);
    }

    if (!empty($tool_calls_trace)) {
      $payload['tool_calls_trace'] = $this->normalizeToolCallsTrace($tool_calls_trace);
    }

    $persisted_renderables = apply_filters('kimaai_chatbot_persisted_renderables', [], $provider, $payload, $response);

    if (!empty($persisted_renderables) && is_array($persisted_renderables)) {
      $response['_kimaai_renderables'] = Helper::normalizeToolTraceValue($persisted_renderables);
    }

    // Unset unnecessary fields from the payload to avoid saving them in the database.
    // In OpenAI, the user_message is redundant since the input messages are already saved.
    if (in_array($provider, ['openai', 'openrouter', 'ollama'], true)) {
      unset($payload['user_message']);
    }
    unset(
      $payload['ai_provider_object'],
      $payload['model'],
      $payload['input'],
    );

    // Save the chat history.
    $chat_data = [
      'user_id' => get_current_user_id(),
      'ip' => $payload['is_playground'] ? 0 : $ip_address,
      'session_id' => $payload['session_id'],
      'input' => wp_json_encode($payload),
      'response' => wp_json_encode($response),
    ];

    $this->chat_model->insert($chat_data);

    $return = [
      'message' => $ai_response,
      'session_id' => $payload['session_id'],
    ];

    if (!empty($response['_kimaai_renderables']) && is_array($response['_kimaai_renderables'])) {
      $return['renderables'] = apply_filters(
        'kimaai_chatbot_hydrated_renderables',
        [],
        $response['_kimaai_renderables'],
        $response,
        $payload,
        []
      );
    }

    return $return;
  }

  /**
   * Save the image response to the database.
   *
   * @param string $provider 'openai' or 'replicate'.
   * @param string $model Model name.
   * @param array $payload The payload data.
   * @param array $response The response data.
   *
   * @return array The saved response data.
   */
  public function saveImageResponse(string $provider, string $model, array $payload, array $response): array {
    $items = $response['output'] ?? $response['data'] ?? [];
    $saved_urls = [];

    foreach ($items as $item) {
      $saved = null;

      // Case 1: plain string → assume it's a URL (Replicate often returns URLs).
      if (is_string($item)) {
        $saved = $this->image_model->saveRemoteImage($item);
      }
      // Case 2: structured array from providers.
      elseif (is_array($item)) {
        // OpenAI Images API can return either 'url' or 'b64_json'.
        if (!empty($item['url']) && is_string($item['url'])) {
          $saved = $this->image_model->saveRemoteImage($item['url']);
        }
        elseif (!empty($item['b64_json']) && is_string($item['b64_json'])) {
          // Default mime for OpenAI b64 is PNG.
          $saved = $this->image_model->saveBase64Image($item['b64_json'], 'image/png');
        }
      }

      if ($saved) {
        $saved_urls[] = $saved;
      }
    }

    $ip_address = isset($_SERVER['REMOTE_ADDR'])
    ? filter_var(wp_unslash($_SERVER['REMOTE_ADDR']), FILTER_VALIDATE_IP)
    : '';

    $session_id = $payload['is_playground'] ? 0 : ($payload['session_id'] ?? 0);
    $is_playground = $payload['is_playground'] ?? false;

    // keep provider/model in allowed_image_args for auditability.
    $provider_object = $payload['ai_provider_object'] ?? ($payload['aiProviderObject'] ?? null);

    $allowed_image_args = [];

    if ($provider_object && property_exists($provider_object, 'allowed_image_args')) {
      // make sure it's an array, then extend with audit fields
      $provider_object->allowed_image_args = array_merge(
        (array) $provider_object->allowed_image_args,
        [
          // value doesn't matter here – we only test keys later
          'ai_provider' => 'string',
          'ai_model' => 'string',
        ]
      );

      $allowed_image_args = $provider_object->allowed_image_args;
    }

    if ($allowed_image_args) {
      $payload_to_save = array_filter(
        $payload,
        static function ($key) use ($allowed_image_args) {
          return isset($allowed_image_args[$key]);
        },
        ARRAY_FILTER_USE_KEY
      );
    }
    else {
      // Fallback: save only a minimal set if for some reason we don't have allowed_image_args
      $payload_to_save = [
        'prompt' => $payload['prompt'] ?? '',
        'ai_provider' => $provider,
        'ai_model' => $model,
      ];
    }

    // In OpenAI, the model is not needed in the payload to be saved due to the ai_model already being in the payload.
    unset(
      $payload['ai_provider_object'],
      $payload['allowed_image_args'],
      $payload['ai_image_models'],
      $payload['env'],
      $payload['input']
    );

    $this->image_model->insert(
      [
        'user_id' => get_current_user_id(),
        'ip' => $is_playground ? 0 : $ip_address,
        'session_id' => $session_id,
        'input' => $payload_to_save,
        'image_paths' => $saved_urls,
      ]
    );

    return [
      'images' => $saved_urls,
    ];
  }

  /**
   * Handles response parsing for different AI providers and models.
   *
   * @param string $provider The provider name.
   * @param array $response The response data from the AI provider.
   *
   * @return string The parsed response.
   *
   * @throws RuntimeException When the provider is unknown.
   */
  public function parseResponse(string $provider, array $response): string {
    if ($provider === 'replicate') {
      return trim(implode(' ', $response['output'] ?? []));
    }

    if (in_array($provider, ['openai', 'openrouter', 'ollama'], true)) {
      // Responses API shortcut.
      if (!empty($response['output_text']) && is_string($response['output_text'])) {
        return $response['output_text'];
      }

      // Responses API structured output.
      $output = $response['output'] ?? [];
      if (is_array($output)) {
        for ($i = count($output) - 1; $i >= 0; $i--) {
          $item = $output[$i];
          if (!is_array($item)) {
            continue;
          }
          if (($item['type'] ?? '') !== 'message') {
            continue;
          }
          if (($item['role'] ?? '') !== 'assistant') {
            continue;
          }

          $text = '';
          foreach (($item['content'] ?? []) as $part) {
            if (is_array($part) && ($part['type'] ?? '') === 'output_text' && isset($part['text'])) {
              $text .= (string) $part['text'];
            }
          }

          if ($text !== '') {
            return $text;
          }
        }
      }

      // Chat Completions fallback.
      $content = $response['choices'][0]['message']['content'] ?? null;

      if (is_string($content)) {
        return $content;
      }

      if (is_array($content)) {
        $text = '';
        foreach ($content as $part) {
          if (!is_array($part)) {
            continue;
          }

          $type = (string) ($part['type'] ?? '');
          if (($type === 'text' || $type === 'output_text') && isset($part['text'])) {
            $text .= (string) $part['text'];
          }
        }

        return $text;
      }

      return '';
    }

    throw new RuntimeException('Unknown AI provider: ' . esc_html($provider));
  }

  /**
   * Checks the connection to the AI service.
   *
   * @param string $provider The AI provider (e.g., 'openai', 'replicate').
   * @param string $api_key The API key for the AI provider.
   *
   * @return bool Returns true if the connection is successful, false otherwise.
   */
  public function testConnection(string $provider, string $api_key, array $options = []): bool {
    try {
      // Local Ollama generally does not require auth. Use a harmless placeholder
      if ($provider === 'ollama' && empty($api_key)) {
        $api_key = 'ollama';
      }

      $client = Factory::create($provider, $api_key, null, $this->normalizeProviderOptions($provider, $options));

      $ids = [];

      // New provider-layer-first path.
      if ($client instanceof ModelClientInterface) {
        foreach ($client->listModels() as $model) {
          if (is_object($model) && !empty($model->id)) {
            $ids[] = (string) $model->id;
          }
        }
      }

      Logger::info(
        "Connection check successful for provider $provider. Available models: " .
        (!empty($ids) ? implode(', ', $ids) : 'none returned')
      );

      return true;
    } catch (\Throwable $exception) {
      Logger::error("Connection check failed for provider $provider: " . $exception->getMessage());
      return false;
    }
  }

  /**
   * Extract tool/function-call trace from provider response metadata.
   */
  private function extractToolCallsTraceFromResponse(string $provider, array $response): array {
    if (!in_array($provider, ['openai', 'openrouter', 'ollama'], true)) {
      return [];
    }

    $trace = $response['_kimaai_tool_calls'] ?? $response['kimaai_tool_calls'] ?? [];

    return is_array($trace) ? $trace : [];
  }

  /**
   * Normalize tool trace into a JSON-safe array.
   */
  private function normalizeToolCallsTrace(array $trace): array {
    $out = [];

    foreach ($trace as $row) {
      if (!is_array($row)) {
        continue;
      }
      $out[] = Helper::normalizeToolTraceValue($row);
    }

    return array_values($out);
  }

  /**
   * Normalize provider-specific connection options from resolved env data.
   */
  private function providerOptionsFromEnv(array $env): array {
    return $this->normalizeProviderOptions(
      (string) ($env['provider'] ?? ''),
      [
        'baseUrl' => !empty($env['endpoint']) && is_string($env['endpoint'])
          ? trim($env['endpoint'])
          : '',
      ]
    );
  }

  /**
   * Apply preset-level generation params (temperature / reasoning / verbosity)
   * to the input payload, but only when the env was resolved from a logical
   * alias and the caller has not already set these keys.
   */
  private function applyPresetParams(array $env, string $provider, string $model, array $input): array {
    $preset_params = $env['preset_params'] ?? [];
    if (empty($preset_params)) {
      return $input;
    }

    $model_def = AiModels::getModel($provider, $model);
    $params = (is_array($model_def) && isset($model_def['params']) && is_array($model_def['params'])) ? $model_def['params'] : [];
    $supports_reasoning = !empty($params['reasoning']);
    $supports_verbosity = !empty($params['verbosity']);

    if ($supports_reasoning || $supports_verbosity) {
      if ($supports_reasoning && !empty($preset_params['reasoning']) && !isset($input['reasoning'])) {
        $input['reasoning'] = ['effort' => (string) $preset_params['reasoning']];
      }
      if ($supports_verbosity && !empty($preset_params['verbosity']) && !isset($input['text'])) {
        $input['text'] = ['verbosity' => (string) $preset_params['verbosity']];
      }
    } elseif (!isset($input['temperature']) && isset($preset_params['temperature'])) {
      $input['temperature'] = (float) $preset_params['temperature'];
    }

    return $input;
  }

  /**
   * Normalize provider-specific connection options.
   */
  private function normalizeProviderOptions(string $provider, array $options): array {
    $provider = strtolower(trim($provider));

    $base_url = isset($options['baseUrl']) && is_string($options['baseUrl'])
    ? trim($options['baseUrl'])
    : '';

    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, '/')];
  }
}
