<?php

namespace KupShop\LLMBundle\Dto;

class PromptResponseStats
{
    private int $totalTokens = 0;
    private int $inputTokens = 0;
    private int $outputTokens = 0;
    private string $model = 'unknown';

    public function setTotalTokens(?int $totalTokens): PromptResponseStats
    {
        $this->totalTokens = $totalTokens ?? 0;

        return $this;
    }

    public function setInputTokens(?int $inputTokens): PromptResponseStats
    {
        $this->inputTokens = $inputTokens ?? 0;

        return $this;
    }

    public function setOutputTokens(?int $outputTokens): PromptResponseStats
    {
        $this->outputTokens = $outputTokens ?? 0;

        return $this;
    }

    public function setModel(?string $model): PromptResponseStats
    {
        $this->model = $model ?? $this->model;

        return $this;
    }

    public function getTotalTokens(): int
    {
        return $this->totalTokens;
    }

    public function getInputTokens(): int
    {
        return $this->inputTokens;
    }

    public function getOutputTokens(): int
    {
        return $this->outputTokens;
    }

    public function getLabels(): array
    {
        return [
            'model' => $this->model ?? null,
        ];
    }

    public function toArray(): array
    {
        return [
            'totalTokens' => $this->totalTokens,
            'inputTokens' => $this->inputTokens,
            'outputTokens' => $this->outputTokens,
            'model' => $this->model,
        ];
    }
}
