<?php

declare(strict_types=1);

namespace KimaAI\Ai\Providers\OpenAI\Stream;

use KimaAI\Ai\Providers\DTO\StreamChunk;
use KimaAI\Ai\Providers\Support\SseParser;

final class ChatCompletionsStreamParser {
  /**
   * Parser for SSE frames.
   *
   * @var SseParser
   */
  private SseParser $parser;

  public function __construct(?SseParser $parser = null) {
    $this->parser = $parser ?? new SseParser();
  }

  /**
   * Parse streaming chat completion chunks.
   *
   * @return list<StreamChunk>
   */
  public function push(string $chunk): array {
    $chunks = [];

    foreach ($this->parser->push($chunk) as $frame) {
      $data = trim($frame['data']);
      if ($data === '' || $data === '[DONE]') {
        continue;
      }

      $json = json_decode($data, true);
      if (!is_array($json)) {
        $chunks[] = new StreamChunk(StreamChunk::TYPE_RAW, null, ['data' => $data], []);
        continue;
      }

      foreach (($json['choices'] ?? []) as $choice) {
        if (!is_array($choice)) {
          continue;
        }

        $delta = $choice['delta'] ?? [];
        if (is_array($delta) && isset($delta['content']) && $delta['content'] !== '') {
          $chunks[] = new StreamChunk(
            StreamChunk::TYPE_TEXT_DELTA,
            (string) $delta['content'],
            [],
            $json
          );
        }

        if (($choice['finish_reason'] ?? null) !== null) {
          $chunks[] = new StreamChunk(
            StreamChunk::TYPE_COMPLETED,
            null,
            ['finish_reason' => $choice['finish_reason']],
            $json
          );
        }
      }

      if ($chunks === [] || end($chunks)->raw !== $json) {
        $chunks[] = new StreamChunk(StreamChunk::TYPE_RAW, null, ['object' => $json['object'] ?? ''], $json);
      }
    }

    return $chunks;
  }
}
