<?php

declare(strict_types=1);

namespace KimaAI\Ai\Providers\Replicate;

use KimaAI\Ai\Providers\AbstractClient;
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\ChatResponse;
use KimaAI\Ai\Providers\DTO\EmbeddingRequest;
use KimaAI\Ai\Providers\DTO\EmbeddingResponse;
use KimaAI\Ai\Providers\DTO\ImageRequest;
use KimaAI\Ai\Providers\DTO\ImageResponse;
use KimaAI\Ai\Providers\DTO\ModelInfo;
use KimaAI\Ai\Providers\Http\HttpClientInterface;
use KimaAI\Ai\Providers\Support\Capability;
use KimaAI\Ai\Providers\Support\ProviderMetadata;


final class ReplicateClient extends AbstractClient implements ChatClientInterface, ImageClientInterface, EmbeddingClientInterface, ModelClientInterface {

  /**
   * Allowed AI API arguments.
   *
   * @var array
   */
  public array $allowed_args = [
    'prompt' => 'string',
    'max_tokens' => 'int',
    'min_tokens' => 'int',
    'temperature' => 'float',
    'top_p' => 'float',
    'top_k' => 'int',
    'stop_sequences' => 'string', // <|end_of_text|>
    'presence_penalty' => 'float',
    'frequency_penalty' => 'float',
    'length_penalty' => 'float',
    'seed' => 'int',
    'prompt_template' => 'string',
    'log_performance_metrics' => 'bool',
  ];

  /**
   * Allowed AI API arguments for image generation.
   *
   * @var array
   */
  public array $allowed_image_args = [
    'prompt' => 'string',
    'num_outputs' => 'int',
    'output_format' => 'string', // png, webp, jpg.
    'aspect_ratio' => 'string', // phpcs:ignore 1:1, 16:9, 4:3.
    'num_inference_steps' => 'int',
    'output_quality' => 'int', // 1-100
    'disable_safety' => 'bool',
    'go_fast' => 'bool',
  ];

  /**
   * Allowed embedding arguments.
   *
   * @var array
   */
  public array $allowed_embedding_args = [
    'text' => 'string',
    'texts' => 'string',
    'path' => 'string',
    'batch_size' => 'int',
    'convert_to_numpy' => 'bool',
    'normalize_embeddings' => 'bool',
  ];

  public function __construct(
    string $api_key,
    string $base_url = 'https://api.replicate.com/v1',
    ?HttpClientInterface $http_client = null
  ) {
    parent::__construct($api_key, $base_url, 300, $http_client);
  }

  public function metadata(): ProviderMetadata {
    return new ProviderMetadata('replicate', 'Replicate', $this->base_url, [
      Capability::CHAT,
      Capability::IMAGE,
      Capability::EMBEDDING,
      Capability::MODELS,
      Capability::REPLICATE_PREDICTIONS,
    ]);
  }

  protected function buildHeaders(bool $stream): array {
    return [
      'Authorization: Bearer ' . $this->api_key,
      'Content-Type: application/json',
    ];
  }

  /**
   * Handles the chat request.
   *
   * @param array $payload The request data.
   * @return array The response data.
   */
  public function chat(array $payload): array {
    $model = $payload['model'];
    $response = $this->predictions()->createForModel($model, $payload);
    return $response;
  }

  public function createChat(ChatRequest $request): ChatResponse {
    $response = $this->chat($request->toArray());
    return new ChatResponse(
      trim(implode(' ', array_map('strval', $response['output'] ?? []))),
      'replicate',
      $request->model,
      $response
    );
  }

  /**
   * Handles the image request.
   *
   * @param array $payload The request data.
   * @return array The response data.
   */
  public function image(array $payload): array {
    $model = $payload['model'];
    $response = $this->predictions()->createForModel($model, $payload);
    return $response;
  }

  public function createImage(ImageRequest $request): ImageResponse {
    $response = $this->image($request->toArray());
    return new ImageResponse($response['output'] ?? [], $response);
  }

  /**
   * Handles the embeddings request.
   *
   * @param array $payload The request data.
   * @return array The response data.
   */
  public function embedding(array $payload): array {
    $model = $payload['model'];
    $input = $this->prepareEmbeddingInput($payload);
    $version = $this->resolveEmbeddingModelVersion($model);
    $response = $version
    ? $this->predictions()->createEmbeddingForVersion($version, $input)
    : $this->predictions()->createEmbeddingForModel($model, $input);
    $response = $this->waitForPrediction($response);
    return $this->normalizeEmbeddingResponse($response);
  }

  public function createEmbedding(EmbeddingRequest $request): EmbeddingResponse {
    $response = $this->embedding($request->toArray());
    $vectors = [];
    foreach (($response['data'] ?? []) as $item) {
      if (is_array($item) && isset($item['embedding']) && is_array($item['embedding'])) {
        $vectors[] = $item['embedding'];
      }
    }
    return new EmbeddingResponse($vectors, $response);
  }

  public function listModels(): array {
    $response = $this->models()->list();
    $models = [];

    foreach (($response['results'] ?? []) as $item) {
      if (is_array($item)) {
        $owner = (string) ($item['owner'] ?? '');
        $name = (string) ($item['name'] ?? '');
        $id = ($owner !== '' && $name !== '') ? "{$owner}/{$name}" : ($name ?: $owner);
        if ($id !== '') {
          $models[] = new ModelInfo($id, 'replicate', $item);
        }
      }
    }

    return $models;
  }

  public function predictions(): Endpoint\Predictions {
    return new Endpoint\Predictions($this);
  }

  public function models(): Endpoint\Models {
    return new Endpoint\Models($this);
  }

  private function prepareEmbeddingInput(array $payload): array {
    $input = array_intersect_key($payload, $this->allowed_embedding_args);

    if (isset($input['texts']) && is_array($input['texts'])) {
      $input['texts'] = wp_json_encode(array_values($input['texts']));
    }

    if (!empty($payload['input']) && empty($input['text']) && empty($input['texts']) && empty($input['path'])) {
      $texts = is_array($payload['input']) ? $payload['input'] : [(string) $payload['input']];
      $input['texts'] = wp_json_encode(array_values($texts));
    }

    return $input;
  }

  private function resolveEmbeddingModelVersion(string &$model): ?string {
    if (str_contains($model, ':')) {
      [$model, $version] = explode(':', $model, 2);
      return $version !== '' ? $version : null;
    }

    return $this->fetchLatestVersionId($model);
  }

  private function fetchLatestVersionId(string $model): ?string {
    $parts = explode('/', $model, 2);
    if (count($parts) !== 2 || $parts[0] === '' || $parts[1] === '') {
      return null;
    }

    $meta = $this->models()->get($parts[0], $parts[1]);
    $latest = $meta['latest_version'] ?? null;

    if (is_array($latest)) {
      return (string) ($latest['id'] ?? '');
    }

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

    return null;
  }

  private function normalizeEmbeddingResponse(array $response): array {
    $vectors = $this->extractEmbeddingVectors($response['output'] ?? null);

    if (!$vectors) {
      return $response;
    }

    $response['data'] = array_map(
      static fn(array $vector, int $index): array => [
        'object' => 'embedding',
        'embedding' => $vector,
        'index' => $index,
      ],
      $vectors,
      array_keys($vectors)
    );

    return $response;
  }

  private function waitForPrediction(array $prediction): array {
    $status = (string) ($prediction['status'] ?? '');
    if (!in_array($status, ['starting', 'processing'], true)) {
      return $prediction;
    }

    $prediction_id = (string) ($prediction['id'] ?? '');
    if ($prediction_id === '') {
      return $prediction;
    }

    $started = time();
    do {
      sleep(2);
      $prediction = $this->predictions()->get($prediction_id);
      $status = (string) ($prediction['status'] ?? '');

      if ($status === 'failed') {
        throw new \RuntimeException('Replicate prediction failed: ' . (string) ($prediction['error'] ?? 'Unknown error'));
      }

      if (in_array($status, ['succeeded', 'successful', 'canceled'], true)) {
        return $prediction;
      }
    } while (time() - $started < $this->timeout);

    return $prediction;
  }

  private function extractEmbeddingVectors($output): array {
    if (!is_array($output)) {
      return [];
    }

    if ($this->isVector($output)) {
      return [$output];
    }

    $vectors = [];
    foreach ($output as $item) {
      if (is_array($item) && $this->isVector($item)) {
        $vectors[] = $item;
      }
    }

    return $vectors;
  }

  private function isVector(array $value): bool {
    if ($value === []) {
      return false;
    }

    foreach ($value as $item) {
      if (!is_int($item) && !is_float($item)) {
        return false;
      }
    }

    return true;
  }
}
