<?php
namespace boru\boruai\Tiktoken\Vocab\Loader;

use RuntimeException;
use boru\boruai\Tiktoken\Vocab\Vocab;
use boru\boruai\Tiktoken\Vocab\VocabLoader;

class DefaultVocabLoader implements VocabLoader {

    /** @var string|null */
    private $cacheDir;

    public function __construct($cacheDir = null) {
        $this->cacheDir = $cacheDir;
    }

    /**
     * @param string $uri
     * @param string|null $checksum
     * @return Vocab
     */
    public function load($uri, $checksum = null) {
        $cacheFile = $this->cacheDir !== null ? $this->cacheDir . DIRECTORY_SEPARATOR . sha1($uri) : null;

        if ($cacheFile !== null) {
            if (file_exists($cacheFile) && $this->checkHash($cacheFile, $checksum)) {
                return Vocab::fromFile($cacheFile);
            }

            assert($this->cacheDir !== null);

            if (! is_dir($this->cacheDir) && ! @mkdir($this->cacheDir, 0750, true)) {
                throw new RuntimeException(sprintf(
                    'Directory does not exist and cannot be created: %s',
                    $this->cacheDir
                ));
            }

            if (! is_writable($this->cacheDir)) {
                throw new RuntimeException(sprintf('Directory is not writable: %s', $this->cacheDir));
            }
        }

        $stream = fopen($uri, 'r');

        if ($stream === false) {
            throw new RuntimeException(sprintf('Could not open stream for URI: %s', $uri));
        }

        try {
            if ($checksum !== null && $this->isRewindable($stream)) {
                if (! $this->checkHash($stream, $checksum)) {
                    throw new RuntimeException(sprintf(
                        'Checksum failed. Could not load vocab from URI: %s',
                        $uri
                    ));
                }

                rewind($stream);
            }

            if ($cacheFile !== null) {
                $cacheStream = fopen($cacheFile, 'w+');

                if ($cacheStream === false) {
                    throw new RuntimeException(sprintf('Could not open file for write: %s', $cacheFile));
                }

                try {
                    stream_copy_to_stream($stream, $cacheStream);

                    return Vocab::fromStream($cacheStream);
                } finally {
                    fclose($cacheStream);
                }
            }

            return Vocab::fromStream($stream);
        } finally {
            fclose($stream);
        }
    }

    /**
     * @param string|resource $resource
     * @param string|null $expectedHash
     * @return bool
     */
    private function checkHash($resource, $expectedHash=null) {
        if ($expectedHash === null) {
            return true;
        }

        $ctx = hash_init('sha256');

        if (is_resource($resource)) {
            hash_update_stream($ctx, $resource);
        } else {
            hash_update_file($ctx, $resource);
        }

        $hash = hash_final($ctx);

        return hash_equals($hash, $expectedHash);
    }

    /**
     * @param resource $stream
     * @return bool
     */
    private function isRewindable($stream) {
        $meta = stream_get_meta_data($stream);

        return $meta['seekable'];
    }
}
