<?php

declare(strict_types=1);

namespace NeuronAI\Providers;

use NeuronAI\Chat\Messages\ContentBlocks\ContentBlockInterface;
use NeuronAI\Chat\Messages\Usage;
use NeuronAI\UniqueIdGenerator;

use function array_values;

class BasicStreamState
{
    protected string $messageId;

    protected array $toolCalls = [];

    /**
     * @var array<string|int, ContentBlockInterface>
     */
    protected array $blocks = [];

    /**
     * @var array<string, mixed>
     */
    protected array $metadata = [];

    public function __construct(
        protected Usage $usage = new Usage(0, 0),
    ) {
    }

    public function addInputTokens(int $tokens): self
    {
        $this->usage->inputTokens += $tokens;
        return $this;
    }

    public function addOutputTokens(int $tokens): self
    {
        $this->usage->outputTokens += $tokens;
        return $this;
    }

    public function getUsage(): Usage
    {
        return $this->usage;
    }

    public function messageId(?string $id = null): string
    {
        if ($id !== null) {
            $this->messageId = $id;
        }

        if (!isset($this->messageId)) {
            $this->messageId = UniqueIdGenerator::generateId('msg_');
        }

        return $this->messageId;
    }

    /**
     * @return ContentBlockInterface[]
     */
    public function getContentBlocks(): array
    {
        return array_values($this->blocks);
    }

    public function hasToolCalls(): bool
    {
        return $this->toolCalls !== [];
    }

    public function getToolCalls(): array
    {
        return $this->toolCalls;
    }

    public function addMetadata(string $key, mixed $value): self
    {
        $this->metadata[$key] = $value;
        return $this;
    }

    public function accumulateMetadata(string $key, string $value): self
    {
        if (!isset($this->metadata[$key])) {
            $this->metadata[$key] = '';
        }

        $this->metadata[$key] .= $value;
        return $this;
    }

    public function getMetadata(?string $key = null): mixed
    {
        if ($key === null) {
            return $this->metadata;
        }

        return $this->metadata[$key] ?? null;
    }

    public function hasMetadata(string $key): bool
    {
        return isset($this->metadata[$key]);
    }
}
