<?php
declare(strict_types=1);

namespace KimaAI\Ai\Providers\OpenAI;

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\Contracts\ResponsesClientInterface;
use KimaAI\Ai\Providers\Contracts\StreamingChatClientInterface;
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\DTO\Usage;
use KimaAI\Ai\Providers\Http\HttpClientInterface;
use KimaAI\Ai\Providers\Http\HttpRequest;
use KimaAI\Ai\Providers\OpenAI\Stream\ChatCompletionsStreamParser;
use KimaAI\Ai\Providers\OpenAI\Stream\ResponsesStreamParser;
use KimaAI\Ai\Providers\Support\Capability;
use KimaAI\Ai\Providers\Support\ProviderMetadata;


class OpenAiClient extends AbstractClient implements ChatClientInterface, StreamingChatClientInterface, ImageClientInterface, EmbeddingClientInterface, ModelClientInterface, ResponsesClientInterface {
  /**
   * Optional OpenAI organization header.
   *
   * @var string|null
   */
  protected ?string $organization;

  /**
   * Allowed AI API arguments (Responses API).
   *
   * @var array
   */
  public array $allowed_args = [
    'model' => 'string',
    'input' => 'array', // []
    'text' => 'object', // {}
    'reasoning' => 'object', // {}
    'tools' => 'array', // []
    'tool_choice' => 'string', // phpcs:ignore (e.g. "auto" | "none" | ["type"=>"function","name"=>"..."])
    'parallel_tool_calls' => 'bool',
    'temperature' => 'float',
    'max_output_tokens' => 'int',
    'top_p' => 'float',
    'store' => 'bool',
    'metadata' => 'array',
    'previous_response_id' => 'string',
    'stream' => 'bool',
    'stream_options' => 'object',

    // Also useful for OpenAI-compatible providers like Ollama.
    'messages' => 'array',
    'instructions' => 'string',
    'truncation' => 'string',
    'frequency_penalty' => 'float',
    'presence_penalty' => 'float',
    'response_format' => 'object',
    'seed' => 'int',
    'stop' => 'array',
    'max_tokens' => 'int',
    'reasoning_effort' => 'string',
  ];

  /**
   * Allowed AI image API arguments.
   *
   * @var array
   */
  public array $allowed_image_args = [
    'model' => 'string',
    'prompt' => 'string',
    'n' => 'int',
    'size' => 'string',
    'output_compression' => 'int',
    'output_format' => 'string',

    // OpenAI-compatible/Ollama-compatible fields.
    'response_format' => 'string',
    'quality' => 'string',
    'style' => 'string',
    'user' => 'string',
  ];

  /**
   * Allowed AI embedding API arguments.
   *
   * @var array
   */
  public array $allowed_embedding_args = [
    'input' => 'string',
    'model' => 'string',
    'encoding_format' => 'string',
    'dimensions' => 'int',
    'user' => 'string',
  ];

  /**
   * AI image models
   *
   * @var array
   */
  public array $ai_image_models = [];

  public function __construct(
    string $api_key,
    ?string $organization = null,
    string $base_url = 'https://api.openai.com/v1',
    ?HttpClientInterface $http_client = null
  ) {
    parent::__construct($api_key, $base_url, 300, $http_client);
    $this->organization = $organization;
    $this->ai_image_models = [
      'dall-e-2' => function_exists('__') ? __('DALL-E 2', 'kimaai') : 'DALL-E 2',
      'dall-e-3' => function_exists('__') ? __('DALL-E 3', 'kimaai') : 'DALL-E 3',
      'gpt-image-1' => function_exists('__') ? __('GPT Image 1', 'kimaai') : 'GPT Image 1',
    ];
  }

  public function metadata(): ProviderMetadata {
    return new ProviderMetadata('openai', 'OpenAI', $this->base_url, [
      Capability::CHAT,
      Capability::STREAMING_CHAT,
      Capability::RESPONSES,
      Capability::IMAGE,
      Capability::EMBEDDING,
      Capability::MODELS,
      Capability::TOOL_CALLING,
      Capability::OPENAI_COMPATIBLE,
    ]);
  }

  public function createChat(ChatRequest $request): ChatResponse {
    $response = $this->chat($request->toArray());

    return new ChatResponse(
      $this->extractText($response),
      $this->metadata()->id,
      (string) ($response['model'] ?? $request->model),
      $response,
      $this->extractUsage($response)
    );
  }

  public function streamChat(ChatRequest $request): iterable {
    $payload = array_intersect_key($request->toArray(), $this->allowed_args);
    $payload['stream'] = true;
    $path = isset($payload['input']) ? '/responses' : '/chat/completions';
    $parser = $path === '/responses'
    ? new ResponsesStreamParser()
    : new ChatCompletionsStreamParser();

    $headers = $this->buildHeaders(true);
    $url = $this->base_url . $path;

    foreach ($this->http_client->stream(new HttpRequest('POST', $url, $headers, $payload, $this->timeout)) as $raw_chunk) {
      foreach ($parser->push($raw_chunk) as $chunk) {
        yield $chunk;
      }
    }
  }

  public function createImage(ImageRequest $request): ImageResponse {
    $response = $this->image($request->toArray());
    return new ImageResponse($response['data'] ?? $response['output'] ?? [], $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['data'] ?? []) as $item) {
      if (is_array($item) && !empty($item['id'])) {
        $models[] = new ModelInfo((string) $item['id'], $this->metadata()->id, $item);
      }
    }

    return $models;
  }

  /**
   * Handles the chat request (non-streaming).
   */
  public function chat(array $payload): array {
    $allowed_args = array_intersect_key($payload, $this->allowed_args);

    if (isset($payload['input'])) {
      return $this->responses()->create($allowed_args, false);
    }

    // Back-compat path if you use another endpoint internally
    return $this->runChat()->create($allowed_args);
  }

  /**
   * Handles the chat request (streaming, Responses API).
   *
   * @param array $payload
   * @param callable $on_chunk function(string $chunk): void
   * @return array
   */
  public function chatStream(array $payload, callable $on_chunk): array {
    $allowed_args = array_intersect_key($payload, $this->allowed_args);

    if (isset($payload['input'])) {
      $allowed_args['stream'] = true;
      return $this->responses()->create($allowed_args, true, $on_chunk);
    }

    $allowed_args['stream'] = true;
    return $this->request('POST', '/chat/completions', $allowed_args, true, [], $on_chunk);
  }

  public function image(array $payload): array {
    $allowed_args = array_intersect_key($payload, $this->allowed_image_args);
    return $this->images()->create($allowed_args);
  }

  public function embedding(array $payload): array {
    $allowed_args = array_intersect_key($payload, $this->allowed_embedding_args);
    return $this->embeddings()->create($allowed_args);
  }

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

    if ($this->organization) {
      $headers[] = 'OpenAI-Organization: ' . $this->organization;
    }

    if ($stream) {
      // OpenAI uses SSE for streaming responses.
      $headers[] = 'Accept: text/event-stream';
    }

    return $headers;
  }

  public function responses(): Endpoint\Responses {
    return new Endpoint\Responses($this);
  }

  public function runChat(): Endpoint\Chat {
    return new Endpoint\Chat($this);
  }

  public function completions(): Endpoint\Completions {
    return new Endpoint\Completions($this);
  }

  public function embeddings(): Endpoint\Embeddings {
    return new Endpoint\Embeddings($this);
  }

  public function images(): Endpoint\Images {
    return new Endpoint\Images($this);
  }

  public function audio(): Endpoint\Audio {
    return new Endpoint\Audio($this);
  }

  public function moderations(): Endpoint\Moderations {
    return new Endpoint\Moderations($this);
  }

  public function assistants(): Endpoint\Assistants {
    return new Endpoint\Assistants($this);
  }

  public function files(): Endpoint\Files {
    return new Endpoint\Files($this);
  }

  public function vectorStores(): Endpoint\VectorStores {
    return new Endpoint\VectorStores($this);
  }

  public function fineTunes(): Endpoint\FineTuning {
    return new Endpoint\FineTuning($this);
  }

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

  protected function extractText(array $response): string {
    if (!empty($response['output_text']) && is_string($response['output_text'])) {
      return $response['output_text'];
    }

    foreach (array_reverse($response['output'] ?? []) as $item) {
      if (!is_array($item) || ($item['type'] ?? '') !== 'message') {
        continue;
      }

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

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

    return (string) ($response['choices'][0]['message']['content'] ?? '');
  }

  protected function extractUsage(array $response): ?Usage {
    $usage = $response['usage'] ?? null;
    if (!is_array($usage)) {
      return null;
    }

    return new Usage(
      isset($usage['input_tokens'])
        ? (int) $usage['input_tokens']
        : (isset($usage['prompt_tokens']) ? (int) $usage['prompt_tokens'] : null),
      isset($usage['output_tokens'])
        ? (int) $usage['output_tokens']
        : (isset($usage['completion_tokens']) ? (int) $usage['completion_tokens'] : null),
      isset($usage['total_tokens']) ? (int) $usage['total_tokens'] : null,
      $usage
    );
  }
}
