<?php

declare(strict_types=1);

namespace KimaAI\Ai\Providers\Replicate\Endpoint;

use InvalidArgumentException;
use KimaAI\Ai\Providers\Replicate\ReplicateClient;


/**
 * Wrapper around Replicate’s predictions endpoints.
 */
final class Predictions {
  /**
   * Replicate client wrapper.
   *
   * @var ReplicateClient
   */
  private ReplicateClient $client;

  /**
   * Text/chat model input spec
   */
  private const ALLOWED_ARGS = [
    'prompt',
    'max_tokens',
    'min_tokens',
    'temperature',
    'top_p',
    'top_k',
    'stop_sequences',
    'presence_penalty',
    'frequency_penalty',
    'length_penalty',
    'seed',
    'prompt_template',
    'log_performance_metrics',
  ];

  /**
   * Image model input spec
   */
  private const ALLOWED_IMAGE_ARGS = [
    'prompt',
    'num_outputs',
    'output_format',
    'aspect_ratio',
    'num_inference_steps',
    'output_quality',
    'disable_safety',
    'go_fast',
  ];

  /**
   * Embedding model input spec
   */
  private const ALLOWED_EMBEDDING_ARGS = [
    'text',
    'texts',
    'path',
    'batch_size',
    'convert_to_numpy',
    'normalize_embeddings',
  ];

  public function __construct(ReplicateClient $client) {
    $this->client = $client;
  }

  /**
   * Synchronous prediction for an *official* Replicate model.
   *
   * OFFICIAL MODEL ENDPOINT
   * POST /v1/models/{owner}/{model}/predictions
   *
   * Example: createForModel(
   *     'meta/llama-2-70b-chat',
   *     ['prompt' => 'Hello world', 'temperature' => 0.7]
   * );
   */
  public function createForModel(string $model_slug, array $input, bool $wait = true): array {
    $input = $this->sanitizeInput($input);
    $headers = $wait ? ['Prefer: wait'] : [];

    return $this->client->request(
      'POST',
      "/models/{$model_slug}/predictions",
      ['input' => $input],
      false,
      $headers
    );
  }

  public function createEmbeddingForModel(string $model_slug, array $input, bool $wait = true): array {
    $input = $this->sanitizeEmbeddingInput($input);
    $headers = $wait ? ['Prefer: wait'] : [];

    return $this->client->request(
      'POST',
      "/models/{$model_slug}/predictions",
      ['input' => $input],
      false,
      $headers
    );
  }

  public function createEmbeddingForVersion(string $version, array $input, bool $wait = true): array {
    $input = $this->sanitizeEmbeddingInput($input);
    $headers = $wait ? ['Prefer: wait'] : [];

    return $this->client->request(
      'POST',
      '/predictions',
      [
        'version' => $version,
        'input' => $input,
      ],
      false,
      $headers
    );
  }

  /**
   * Backwards-compatible call for community model versions.
   *
   * COMMUNITY VERSION ENDPOINT
   * POST /v1/predictions  (version parameter required)
   *
   * $params must contain a 'version' key.
   */
  public function create(array $params): array {
    return $this->client->request('POST', '/predictions', $params);
  }

  /**
   * Fetch or cancel predictions (provided for completeness).
   */
  public function get(string $id): array {
    return $this->client->request('GET', "/predictions/{$id}");
  }

  public function cancel(string $id): array {
    return $this->client->request('POST', "/predictions/{$id}/cancel");
  }

  private function sanitizeInput(array $raw): array {
    $image_only_args = array_diff(self::ALLOWED_IMAGE_ARGS, ['prompt']);
    $is_image = array_intersect_key($raw, array_flip($image_only_args));

    $allowed = $is_image ? self::ALLOWED_IMAGE_ARGS : self::ALLOWED_ARGS;
    $clean = array_intersect_key($raw, array_flip($allowed));

    if (empty($clean['prompt'])) {
      throw new InvalidArgumentException('Missing "prompt" value.');
    }

    return $clean;
  }

  private function sanitizeEmbeddingInput(array $raw): array {
    $clean = array_intersect_key($raw, array_flip(self::ALLOWED_EMBEDDING_ARGS));

    if (empty($clean['text']) && empty($clean['texts']) && empty($clean['path'])) {
      throw new InvalidArgumentException('Missing "text", "texts", or "path" value.');
    }

    return $clean;
  }
}
