<?php

declare(strict_types=1);

namespace ContextPaging;

/**
 * Token counter implementation using the token-counter binary.
 */
class TokenCounter implements OpenAICompletionsInterface
{
    private string $binaryPath;

    public function __construct(?string $binaryPath = null)
    {
        $this->binaryPath = $binaryPath ?? __DIR__ . '/../token-counter';
    }

    /**
     * Count tokens in a string.
     *
     * @param string $text The text to tokenize.
     * @param string $encoding One of: cl100k_base (GPT-4/3.5), o200k_base (GPT-4o/o1).
     * @return int Token count.
     * @throws \RuntimeException on process failure.
     */
    public function count(string $text, string $encoding = 'cl100k_base'): int
    {
        $cmd = escapeshellarg($this->binaryPath) . ' ' . escapeshellarg($encoding);

        $descriptors = [
            0 => ['pipe', 'r'], // stdin - we write the text here
            1 => ['pipe', 'w'], // stdout - we read the count here
            2 => ['pipe', 'w'], // stderr - capture errors
        ];

        $process = proc_open($cmd, $descriptors, $pipes);

        if (!is_resource($process)) {
            throw new \RuntimeException('Failed to start token-counter process');
        }

        // Write text to stdin in chunks to avoid pipe buffer deadlocks on large inputs.
        $this->writeToStdin($pipes[0], $text);
        fclose($pipes[0]);

        // Read stdout (the token count)
        $stdout = stream_get_contents($pipes[1]);
        fclose($pipes[1]);

        // Read stderr (any errors)
        $stderr = stream_get_contents($pipes[2]);
        fclose($pipes[2]);

        $exitCode = proc_close($process);

        if ($exitCode !== 0) {
            throw new \RuntimeException("token-counter failed (exit {$exitCode}): {$stderr}");
        }

        $count = trim($stdout);
        if (!ctype_digit($count)) {
            throw new \RuntimeException("Unexpected output from token-counter: {$stdout}");
        }

        return (int) $count;
    }

    /**
     * Count tokens for an array of messages (chat format).
     * Concatenates all content with separator tokens for a rough count.
     *
     * @param array $messages Array of ['role' => '...', 'content' => '...']
     * @param string $encoding Encoding name.
     * @param int $perMessage Overhead tokens per message (4 for cl100k models).
     * @param int $replyPrimer Tokens added after all messages (3 for cl100k models).
     * @return int
     */
    public function contextSize(
        array $messages,
        string $encoding = 'cl100k_base',
        int $perMessage = 4,
        int $replyPrimer = 3
    ): int {
        // Concatenate all role + content so we only fork once.
        $parts = [];
        foreach ($messages as $msg) {
            $parts[] = ($msg['role'] ?? '') . "\n" . ($msg['content'] ?? '');
        }
        $blob = implode("\n", $parts);

        $rawTokens = $this->count($blob, $encoding);

        // Add per-message overhead + reply primer
        return $rawTokens + (count($messages) * $perMessage) + $replyPrimer;
    }

    /**
     * Write large strings to stdin without deadlocking.
     */
    private function writeToStdin($pipe, string $data): void
    {
        $chunkSize = 8192; // 8KB chunks
        $offset = 0;
        $length = strlen($data);

        while ($offset < $length) {
            $chunk = substr($data, $offset, $chunkSize);
            $written = fwrite($pipe, $chunk);
            if ($written === false) {
                throw new \RuntimeException('Failed to write to token-counter stdin');
            }
            $offset += $written;
        }
    }
}
