<?php

/*
 * This file is part of the Symfony package.
 *
 * (c) Fabien Potencier <fabien@symfony.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

namespace Symfony\AI\Chat;

use Symfony\AI\Platform\Message\AssistantMessage;
use Symfony\AI\Platform\Message\MessageBag;
use Symfony\AI\Platform\Message\UserMessage;
use Symfony\Component\Clock\ClockInterface;
use Symfony\Component\Clock\MonotonicClock;
use Symfony\Contracts\Service\ResetInterface;

/**
 * @author Guillaume Loulier <personal@guillaumeloulier.fr>
 *
 * @phpstan-type ChatData array{
 *      action: string,
 *      bag?: MessageBag,
 *      message?: UserMessage,
 *      submitted_at?: \DateTimeImmutable,
 *      streamed_at?: \DateTimeImmutable,
 *      initiated_at?: \DateTimeImmutable,
 *  }
 */
final class TraceableChat implements ChatInterface, ResetInterface
{
    /**
     * @var ChatData[]
     */
    private array $calls = [];

    public function __construct(
        private readonly ChatInterface $chat,
        private readonly ClockInterface $clock = new MonotonicClock(),
    ) {
    }

    public function initiate(MessageBag $messages): void
    {
        $this->calls[] = [
            'action' => 'initiate',
            'bag' => $messages,
            'initiated_at' => $this->clock->now(),
        ];

        $this->chat->initiate($messages);
    }

    public function submit(UserMessage $message): AssistantMessage
    {
        $this->calls[] = [
            'action' => 'submit',
            'message' => $message,
            'submitted_at' => $this->clock->now(),
        ];

        return $this->chat->submit($message);
    }

    public function stream(UserMessage $message): \Generator
    {
        $this->calls[] = [
            'action' => 'stream',
            'message' => $message,
            'streamed_at' => $this->clock->now(),
        ];

        return $this->chat->stream($message);
    }

    /**
     * @return ChatData[]
     */
    public function getCalls(): array
    {
        return $this->calls;
    }

    public function reset(): void
    {
        if ($this->chat instanceof ResetInterface) {
            $this->chat->reset();
        }

        $this->calls = [];
    }
}
