<?php

declare(strict_types=1);

namespace ContextPaging;

use GuzzleHttp\Psr7\ServerRequest;
use Psr\Http\Message\ServerRequestInterface;

/**
 * Extended ServerRequest with context size tracking.
 */
class ContextRequest extends ServerRequest
{
    private ?TokenCounter $tokenCounter = null;

    /**
     * Create from a standard ServerRequestInterface.
     */
    public static function fromRequest(ServerRequestInterface $request): self
    {
        return new self(
            method: $request->getMethod(),
            uri: $request->getUri(),
            headers: $request->getHeaders(),
            body: $request->getBody(),
            version: $request->getProtocolVersion(),
            serverParams: $request->getServerParams()
        );
    }

    /**
     * Set the token counter instance.
     */
    public function withTokenCounter(TokenCounter $counter): self
    {
        $new = clone $this;
        $new->tokenCounter = $counter;
        return $new;
    }

    /**
     * Get the token counter (lazy-load default if not set).
     */
    public function getTokenCounter(): TokenCounter
    {
        if ($this->tokenCounter === null) {
            $this->tokenCounter = new TokenCounter();
        }
        return $this->tokenCounter;
    }

    /**
     * Calculate and store context size for the given messages.
     *
     * @param array $messages Array of ['role' => '...', 'content' => '...']
     * @param string $encoding Encoding name (cl100k_base or o200k_base).
     * @param int $perMessage Overhead tokens per message.
     * @param int $replyPrimer Tokens added after all messages.
     * @return self New request instance with context_size attribute.
     */
    public function contextSize(
        array $messages,
        string $encoding = 'cl100k_base',
        int $perMessage = 4,
        int $replyPrimer = 3
    ): self {
        $count = $this->getTokenCounter()->contextSize(
            $messages,
            $encoding,
            $perMessage,
            $replyPrimer
        );

        return $this->withAttribute('context_size', [
            'tokens' => $count,
            'encoding' => $encoding,
            'message_count' => count($messages),
            'per_message' => $perMessage,
            'reply_primer' => $replyPrimer,
        ]);
    }

    /**
     * Get the stored context size, if any.
     *
     * @return array|null ['tokens' => int, 'encoding' => string, ...]
     */
    public function getContextSize(): ?array
    {
        return $this->getAttribute('context_size');
    }

    /**
     * Get just the token count from stored context size.
     */
    public function getContextTokenCount(): ?int
    {
        $ctx = $this->getContextSize();
        return $ctx['tokens'] ?? null;
    }
}