<?php

/**
 * trait -> ai
 * 
 * @package Sngine
 * @author Zamblek
 */

trait AiTrait
{

  /* ------------------------------- */
  /* AI */
  /* ------------------------------- */

  /**
   * generate_ai_text
   * 
   * @param string $prompt
   * @param int $max_result_length
   * @param bool $is_blog
   * @return array
   */
  public function generate_ai_text($prompt, $max_result_length, $is_blog = false)
  {
    global $system, $db;
    $max_result_length = $max_result_length ? (int)$max_result_length : 0;
    /* validate prompt */
    if (empty($prompt)) {
      throw new ValidationException(__("You need to enter a prompt to generate text"));
    }
    /* validate max result length */
    if ($max_result_length && (!is_numeric($max_result_length) || $max_result_length <= 0)) {
      throw new ValidationException(__("Max result length must be greater than 0"));
    }
    /* check if credit enabled */
    if ($system['ai_text_credit_enabled']) {
      /* set max result length */
      $user_max_credits = floor($this->_data['user_ai_credits_balance'] / $system['ai_generated_word_cost']);
      $max_result_length = ($max_result_length > 0) ? min($max_result_length, $user_max_credits) : $user_max_credits;
      /* check if user has enough credits */
      if ($this->_data['user_ai_credits_balance'] <= 0 || $this->_data['user_ai_credits_balance'] < $system['ai_generated_word_cost'] || $this->_data['user_ai_credits_balance'] < $system['ai_generated_word_cost'] * $max_result_length) {
        throw new ValidationException(__("You do not have enough AI credits to generate this text"));
      }
    }
    /* generate ai text */
    if ($system['ai_mode'] == 'live') {
      $client = OpenAI::client($system['openai_api_key']);
      /* adjust prompt for title generation */
      if ($is_blog) {
        $prompt = "Generate a title and the following content. Format the response as: \nTitle: [Title Here]\nContent: [Content Here]\n\nPrompt: " . $prompt;
      }
      $payload = [
        'model' => $system['openai_text_model'],
        'input' => $prompt,
      ];
      if ($max_result_length) {
        $payload['max_output_tokens'] = $max_result_length;
      }
      $response = $client->responses()->create($payload);
      $full_response = $response->outputText;
      /* extract title if requested */
      $title = "";
      $generated_text = $full_response;
      if ($is_blog) {
        // Use regex or string splitting to separate Title and Content
        if (preg_match('/Title:\s*(.*?)\nContent:\s*(.*)/is', $full_response, $matches)) {
          $title = trim($matches[1]);
          $generated_text = trim($matches[2]);
        }
      }
    } else {
      $title = __("AI Generated Title");
      $generated_text = __("This is a sample text generated by AI. It is used to test the AI system.");
      $full_response = $title . "\n" . $generated_text;
    }
    /* credit calculation (using full response to be safe) */
    if ($system['ai_text_credit_enabled']) {
      /* calculate generated text credits */
      $generated_text_credits = count_words($full_response) * $system['ai_generated_word_cost'];
      /* decrease user ai credits balance */
      $this->_data['user_ai_credits_balance'] -= $generated_text_credits;
      $db->query(sprintf('UPDATE users SET user_ai_credits_balance = user_ai_credits_balance - %s WHERE user_id = %s', secure($generated_text_credits, 'int'), secure($this->_data['user_id'], 'int')));
    }
    /* return */
    return [
      'generated_title' => $title,
      'generated_text' => $generated_text,
      'ai_credits' => ($system['ai_text_credit_enabled']) ? $this->_data['user_ai_credits_balance'] : 0,
      'words_balance' => ($system['ai_text_credit_enabled']) ? floor($this->_data['user_ai_credits_balance'] / $system['ai_generated_word_cost']) : 0,
    ];
  }


  /**
   * generate_ai_image
   * 
   * @param string $prompt
   * @param string $image_size
   * @return array
   */
  public function generate_ai_image($prompt, $image_size)
  {
    global $system, $db;
    /* validate prompt */
    if (empty($prompt)) {
      throw new ValidationException(__("You need to enter a prompt to generate an image"));
    }
    /* validate image size */
    if (!in_array($image_size, ['auto', '1024x1024', '1024x1536', '1536x1024'])) {
      throw new ValidationException(__("You need to select an image size"));
    }
    /* check if credit enabled */
    if ($system['ai_image_credit_enabled']) {
      /* check if user has enough credits */
      if ($this->_data['user_ai_credits_balance'] <= 0 || $this->_data['user_ai_credits_balance'] < $system['ai_generated_image_cost']) {
        throw new ValidationException(__("You do not have enough AI credits to generate this image"));
      }
    }
    if ($system['ai_mode'] == 'live') {
      switch ($system['ai_images_provider']) {
        case 'openai':
          /* validate openai api key */
          if (is_empty($system['openai_api_key'])) {
            throw new ValidationException(__("You need to enter a valid OpenAI API key"));
          }
          /* generate ai image ⇢ openai */
          $client = OpenAI::client($system['openai_api_key']);
          $payload = [
            'model' => $system['openai_image_model'],
            'prompt' => $prompt,
            'size' => $image_size,
          ];
          $response = $client->images()->create($payload);
          $generated_image = $response->data[0]->b64_json;
          break;

        case 'replicate':
          /* validate replicate api token */
          if (is_empty($system['replicate_api_token'])) {
            throw new ValidationException(__("You need to enter a valid Replicate API token"));
          }
          /* generate ai image ⇢ replicate */
          $ch = curl_init("https://api.replicate.com/v1/models/" . $system['replicate_model'] . "/predictions");
          curl_setopt($ch, CURLOPT_POST, 1);
          curl_setopt($ch, CURLOPT_RETURNTRANSFER, true);
          curl_setopt($ch, CURLOPT_HTTPHEADER, [
            "Authorization: Bearer " . $system['replicate_api_token'],
            "Content-Type: application/json",
            "Prefer: wait"
          ]);
          $input = [
            'prompt' => $prompt,
            'aspect_ratio' => ($image_size == '1024x1024') ? '1:1' : '4:3',
          ];
          if (!is_empty($system['replicate_denoising_steps'])) {
            $input['num_inference_steps'] = (int)$system['replicate_denoising_steps'];
          }
          if (!is_empty($system['replicate_guidance_scale'])) {
            $input['guidance_scale'] = (float)$system['replicate_guidance_scale'];
          }
          if (!is_empty($system['replicate_seed'])) {
            $input['seed'] = (int)$system['replicate_seed'];
          }
          $payload = json_encode(['input' => $input]);
          curl_setopt($ch, CURLOPT_POSTFIELDS, $payload);
          $result = json_decode(curl_exec($ch), true);
          curl_close($ch);
          if (!isset($result['output'])) {
            throw new Exception($result['error'] ?? "Replicate generation failed");
          }
          $image_url = is_array($result['output']) ? $result['output'][0] : $result['output'];
          $generated_image = base64_encode(file_get_contents($image_url));
          break;
      }
    } else {
      $image_url = "https://placehold.co/1024x1024/EEE/31343C/png?text=AI+Generated+Sample+Image&font=poppins";
      $generated_image = base64_encode(file_get_contents($image_url));
    }
    /* decrease user ai credits balance if credit enabled */
    if ($system['ai_image_credit_enabled']) {
      $this->_data['user_ai_credits_balance'] -= $system['ai_generated_image_cost'];
      $db->query(sprintf('UPDATE users SET user_ai_credits_balance = user_ai_credits_balance - %s WHERE user_id = %s', secure($system['ai_generated_image_cost'], 'int'), secure($this->_data['user_id'], 'int')));
    }
    /* return */
    return [
      'generated_image' => $generated_image,
      'ai_credits' => ($system['ai_image_credit_enabled']) ? $this->_data['user_ai_credits_balance'] : 0,
      'images_balance' => ($system['ai_image_credit_enabled']) ? floor($this->_data['user_ai_credits_balance'] / $system['ai_generated_image_cost']) : 0,
    ];
  }
}
