<?php
// SPDX-License-Identifier: GPL-3.0-or-later
declare(strict_types=1);

namespace Aiwend\Knowledge;

use Aiwend\Db\Connection;
use Aiwend\Support\Clock;
use Aiwend\Usage\Ledger;

/**
 * Builds an index generation, and keeps the active one current, in bounded, resumable steps. One step() does a little work and
 * reports what to do next; nothing is held between steps except the data itself, so any step can be the last (crash, lease loss,
 * duplicate job) and the next one continues from the database.
 *
 * Source identity, not counts: the work list is IndexGenerations::missingSources(), i.e. ready public sources with no ready
 * generation version whose content hash equals the source's current version hash. A source that is added, replaced, made public or
 * restored while a build runs shows up (or reappears) in that list on the next step; a source that is deleted or made private drops
 * out of it. A build only becomes "ready" through IndexGenerations::finish(), which re-checks everything at that instant.
 *
 * Version identity: a source is built only from content whose hash equals the hash of its current lexical version. If the stored
 * file or the post has moved on but the lexical re-index has not run yet, the source is DEFERRED (never built from text the single
 * index has not accepted). At publish time the version must still be of the source's current content (checked inside the publishing
 * transaction), so a result for text replaced while it was embedded is never published.
 *
 * Ownership (GenerationClaims): every source is built under a database lease. A second worker finds it busy and moves on; an expired
 * lease (crashed worker) is taken over by exactly one successor; a worker that lost its lease stops at its next batch and cannot
 * publish (the lease is proved inside the publishing transaction). No database lock is held while embedding.
 *
 * Dispatch safety: a successor that takes over a dead worker's task first looks at the ledger for what that worker may have sent
 * (the operation id was announced before the first dispatch). A request still 'reserved' (never sent) is released; one 'dispatched' is
 * marked uncertain; any uncertain request BLOCKS the task until an administrator unblocks it (GenerationClaims::unblock). This prevents
 * duplicate dispatches by cooperating workers; it does not make provider billing exactly-once: a request whose network outcome is
 * unknown stays unknown and held (nothing here releases or settles an uncertain hold), and a rebuild after a crash that happened
 * after a SETTLED batch pays for that batch again, as a known cost.
 *
 * Scheduling: a persistent cursor (knowledge_generations.cursor_source_id) walks the missing sources in id order, a bounded window
 * per step, wrapping to the start at the end of the list. Busy, deferred or blocked sources are passed over for the rest of the cycle,
 * so they cannot starve later sources, and already complete versions are never in the list.
 *
 * Failure policy (a generation is only ever discarded when continuing cannot help):
 *   - a source erased or unreadable: skipped / deferred, never an error;
 *   - untyped local failure (database, time budget): 'retry' (the caller backs off), the source is rebuilt from scratch;
 *   - metered or provider failure (budget, provider, burned operation, no embedder): PAUSED, never retried here; resuming is an
 *     explicit act and uses a new operation id, so a dispatched call is never repeated silently;
 *   - input the provider cannot accept, or a source that cannot be extracted: the build FAILS (vectors deleted at once);
 *   - more chunks than IndexLimits allows for the model's dimension, or a configuration change: the build FAILS before anything is
 *     embedded / abandoned.
 * The active generation is never failed: its failures are noted and the affected source is simply not answered from it (the lexical
 * index is unaffected, and stale text is never served because retrieval requires hash equality).
 *
 * Nothing here contacts a provider by itself, decides prices or enables semantic retrieval; the embedder comes from an
 * EmbedderFactory, which has no production implementation yet.
 */
final class GenerationBuilder
{
    /** Kept for callers that only know the reference dimension; the real ceiling depends on the model's dimension (IndexLimits). */
    public const MAX_CHUNKS = IndexLimits::REFERENCE_CHUNKS;
    /** Per-embedding-call cost allowance per text, same formula as the reservation estimate (bytes + 128). */
    private const PER_TEXT_OVERHEAD = 128;

    private GenerationClaims $claims;
    private IndexLimits $limits;

    /**
     * @param \Closure(array<string, mixed>): (array{raw:string, type:string}|null) $content current text of a source, or null when unavailable
     * @param ?Ledger $ledger only used to inspect and tidy a dead worker's requests (release never-sent ones, mark sent ones uncertain)
     */
    public function __construct(
        private Connection $db,
        private SourceRepository $sources,
        private IndexGenerations $generations,
        private EmbedderFactory $embedders,
        private \Closure $content,
        private Clock $clock,
        private int $maxSources = 5,
        private int $maxChunks = 250,
        private int $timeBudgetSeconds = 25,
        private ?\Closure $afterBatch = null,
        ?GenerationClaims $claims = null,
        ?IndexLimits $limits = null,
        private ?Ledger $ledger = null,
    ) {
        $this->claims = $claims ?? new GenerationClaims($db, $clock);
        $this->limits = $limits ?? new IndexLimits();
    }

    /**
     * What a build for this configuration would involve, before anything is created or embedded (for a future confirmation screen).
     *
     * @return array{ok:bool, code:?string, chunks:int, limit:int, sources:int, units:int}
     */
    public function preflight(int $siteId, EmbeddingConfig $config): array
    {
        $chunks = $this->siteChunks($siteId);
        $dimension = (int) $config->dimension;
        $check = $this->limits->check($dimension, $chunks);
        $row = $this->db->selectOne(
            "SELECT COUNT(DISTINCT s.id) AS n, COALESCE(SUM(LENGTH(c.text)), 0) AS bytes, COUNT(c.id) AS chunks
             FROM {knowledge_sources} s JOIN {knowledge_chunks} c ON c.version_id = s.current_version_id
             WHERE s.site_id = ? AND s.status = 'ready' AND s.visibility = 'public'",
            [$siteId]
        );
        return [
            'ok' => $check['ok'], 'code' => $check['ok'] ? null : ($check['limit'] === 0 ? 'unsupported_dimension' : 'chunk_ceiling'),
            'chunks' => $chunks, 'limit' => $check['limit'], 'sources' => (int) $row['n'],
            'units' => (int) $row['bytes'] + self::PER_TEXT_OVERHEAD * (int) $row['chunks'],
        ];
    }

    /**
     * One bounded unit of work on a BUILDING generation (build) or the ACTIVE one (keep current).
     *
     * @param int $attempt distinguishes embedding operation ids between runs (the queue's attempt number)
     * @param ?EmbeddingConfig $config the site's current configuration; a build whose fingerprint no longer matches is failed
     * @param ?int $tokenCeiling estimate units this build may spend in total; $unitsBefore is what earlier steps already spent
     * @return array{state:string, code:?string, built:int, skipped:int, deferred:int, blocked:int, remaining:int, units:int, source:?int, limit:?int, chunks:?int}
     *   state: continue | waiting | complete | synced | paused | retry | failed | abandoned
     */
    public function step(int $siteId, int $generationId, int $attempt = 1, ?EmbeddingConfig $config = null, ?int $tokenCeiling = null, int $unitsBefore = 0): array
    {
        $gen = $this->generations->find($siteId, $generationId);
        if ($gen === null) {
            return $this->result('abandoned', 'not_found');
        }
        // Privacy first, whatever the state: replaced text must not linger in any generation, retained or not.
        $this->generations->dropStaleVersions($siteId);

        if ($gen['state'] === IndexGenerations::READY && $this->generations->missingSources($siteId, $generationId, 0, 1) !== []) {
            // Complete once, not any more (a source changed while it waited for activation): continue it instead of re-embedding everything.
            try {
                $this->generations->reopen($siteId, $generationId);
            } catch (GenerationException) {
                return $this->result('abandoned', 'not_building'); // activated or ended in the meantime
            }
            $gen = $this->generations->find($siteId, $generationId) ?? $gen;
        }
        $active = $gen['state'] === IndexGenerations::ACTIVE;
        if (!$active && $gen['state'] !== IndexGenerations::BUILDING) {
            return $this->result('abandoned', 'not_building'); // a duplicate or stale job for a finished, failed, superseded or complete generation
        }
        if (!$active && $config !== null && (!$config->isSemantic() || !hash_equals((string) $gen['fingerprint'], (string) $config->fingerprint()))) {
            return $this->giveUp($siteId, $generationId, 'config_changed');
        }
        $this->generations->discardUnclaimed($siteId, $generationId);

        $chunkCount = $this->siteChunks($siteId);
        $check = $this->limits->check((int) $gen['dimension'], $chunkCount);
        if (!$check['ok']) {
            $code = $check['limit'] === 0 ? 'unsupported_dimension' : 'chunk_ceiling';
            $r = $active ? $this->pause($siteId, $generationId, $code) : $this->giveUp($siteId, $generationId, $code);
            return ['limit' => $check['limit'], 'chunks' => $chunkCount] + $r;
        }

        $started = $this->clock->now();
        $out = $this->result('continue', null);
        $embedder = null;
        $attempted = 0;
        $examined = 0;
        $chunks = 0;
        $window = $this->maxSources * 4;
        $cursor = $this->generations->cursor($siteId, $generationId);
        $ids = $this->generations->missingSources($siteId, $generationId, $cursor, $window);
        if ($ids === [] && $cursor > 0) {
            $cursor = 0; // end of the list: start the next cycle
            $ids = $this->generations->missingSources($siteId, $generationId, 0, $window);
        }
        $reachedEnd = count($ids) < $window;
        $lastDone = $cursor;
        foreach ($ids as $sourceId) {
            if ($attempted >= $this->maxSources || ($attempted > 0 && $chunks >= $this->maxChunks) || ($attempted > 0 && $this->clock->now() - $started >= $this->timeBudgetSeconds)) {
                $reachedEnd = false;
                break;
            }
            $examined++;
            $prepared = $this->prepare($siteId, $generationId, $sourceId);
            if (!is_array($prepared)) {
                $out[$prepared]++;
                if ($prepared === 'skipped') {
                    $this->generations->eraseSource($siteId, $sourceId);
                }
                $lastDone = $sourceId;
                continue;
            }
            ['raw' => $raw, 'type' => $type, 'hash' => $hash, 'units' => $units, 'chunks' => $count, 'claim' => $claim] = $prepared;
            if ($tokenCeiling !== null && $unitsBefore + $out['units'] + $units > $tokenCeiling) {
                $this->claims->release($claim);
                return $this->finishStep($siteId, $generationId, $active, $out, $lastDone, 'paused', 'token_ceiling');
            }
            try {
                $embedder ??= $this->embedders->forGeneration($gen);
            } catch (\Throwable) {
                $this->claims->release($claim);
                return $this->finishStep($siteId, $generationId, $active, $out, $lastDone, 'paused', 'embedder_unavailable');
            }
            $attempted++;
            $chunks += $count;
            $operation = sprintf('gen%d:s%d:%s:a%d', $generationId, $sourceId, substr($hash, 0, 12), $attempt);
            $this->claims->announce($claim, $operation, intdiv($count + Ingestor::BATCH - 1, Ingestor::BATCH));
            $beat = function (int $sid, int $batch) use ($claim): void {
                if ($this->afterBatch !== null) {
                    ($this->afterBatch)($sid, $batch);
                }
                if (!$this->claims->renew($claim)) {
                    throw new ClaimLostException(); // ownership moved on: stop before the next batch is sent
                }
            };
            $guard = function () use ($claim): void {
                if (!$this->claims->renew($claim)) {
                    throw new ClaimLostException(); // checked right before each request is sent
                }
            };
            $ingestor = new Ingestor($this->db, $this->sources, new Extractor(), new Chunker(), $embedder, $this->clock, 120, $beat, $operation, beforeBatch: $guard);
            $batches = intdiv($count + Ingestor::BATCH - 1, Ingestor::BATCH);
            try {
                $r = $ingestor->ingestForGeneration($siteId, $sourceId, $raw, $type, $generationId, $active, fn(): bool => $this->claims->renew($claim));
            } catch (ClaimLostException | SourceChangedException $e) {
                // Fenced (a successor owns the task now) or built from text that was replaced meanwhile: nothing was published. Whatever this
                // worker left unfinished is the successor's to clear; a changed source simply shows up again in the work list.
                if ($e instanceof SourceChangedException) {
                    $this->claims->release($claim);
                }
                $out['deferred']++;
                $out['units'] += $units;
                $lastDone = $sourceId;
                continue;
            } catch (GenerationException $e) {
                // The generation ended under us (activated, failed, superseded): stop and change nothing. An embedder that does not fit it is a setup problem.
                $this->claims->release($claim);
                $out['source'] = $sourceId;
                return $e->getMessage() === 'not_building'
                    ? $this->finishStep($siteId, $generationId, $active, $out, $lastDone, 'abandoned', 'not_building')
                    : $this->finishStep($siteId, $generationId, $active, $out, $lastDone, 'paused', 'embedder_mismatch');
            } catch (\Throwable) {
                // Untyped local failure (the ingestor already discarded its partial version). The attempt may have cost something.
                $out['units'] += $units;
                $out['source'] = $sourceId;
                $uncertain = $this->dispatchUncertain($siteId, $operation, $batches);
                $uncertain ? $this->claims->block($claim) : $this->claims->release($claim);
                return $this->finishStep($siteId, $generationId, $active, $out, $lastDone, $uncertain ? 'paused' : 'retry', $uncertain ? 'uncertain_dispatch' : 'source_failed');
            }
            $out['units'] += $units; // every attempt is counted, successful or not: the ceiling is a guard, not an invoice
            if ($r['outcome'] === 'indexed' || $r['outcome'] === 'unchanged') {
                $this->claims->release($claim);
                $out['built']++;
                $lastDone = $sourceId;
                // The source may have changed while it was embedded: a version for outdated text is removed at once.
                $this->generations->dropStaleVersions($siteId, $sourceId);
                continue;
            }
            if ($r['outcome'] === 'skipped') {
                $this->claims->release($claim);
                $out['skipped']++;
                $lastDone = $sourceId;
                $this->generations->eraseSource($siteId, $sourceId);
                continue;
            }
            $out['source'] = $sourceId;
            [$state, $code] = $this->rejection((string) $r['reason']);
            if (in_array($r['reason'], ['embedding_failed', 'embedding_duplicate'], true) && $this->dispatchUncertain($siteId, $operation, $batches)) {
                $this->claims->block($claim);
                [$state, $code] = ['paused', 'uncertain_dispatch'];
            } else {
                $this->claims->release($claim);
            }
            return $this->finishStep($siteId, $generationId, $active, $out, $lastDone, $state, $code);
        }
        return $this->finishStep($siteId, $generationId, $active, $out, $reachedEnd ? 0 : $lastDone, null, null);
    }

    /**
     * Decides whether a source can be built now and, if so, claims it. Returns the work (raw text, type, hash, estimate, claim) or why it is not
     * built now: 'skipped' (erased), 'deferred' (not settled yet: another worker has it, the text is unavailable, or the lexical index has
     * not accepted this content yet) or 'blocked' (an earlier dispatch is uncertain; an administrator must unblock it).
     *
     * @return array{raw:string, type:string, hash:string, units:int, chunks:int, claim:GenerationClaim}|'skipped'|'deferred'|'blocked'
     */
    private function prepare(int $siteId, int $generationId, int $sourceId): array|string
    {
        $source = $this->sources->find($siteId, $sourceId);
        if ($source === null || in_array($source['status'], ['deleting', 'deleted'], true)) {
            return 'skipped';
        }
        if ($source['status'] !== 'ready' || $source['visibility'] !== 'public' || $source['current_version_id'] === null) {
            return 'deferred';
        }
        $current = $this->db->selectOne("SELECT id, content_hash, chunk_count FROM {knowledge_versions} WHERE id = ? AND site_id = ? AND source_id = ? AND state = 'ready'", [$source['current_version_id'], $siteId, $sourceId]);
        $text = ($this->content)($source);
        if ($current === null || $text === null || Ingestor::contentHash($text['type'], $text['raw']) !== $current['content_hash']) {
            return 'deferred';
        }
        $claim = $this->claims->claim($siteId, $generationId, $sourceId, (string) $current['content_hash']);
        if ($claim === 'blocked') {
            return 'blocked';
        }
        if (!$claim instanceof GenerationClaim) {
            return 'deferred'; // another worker is embedding this source right now
        }
        if ($claim->previous !== null && $this->dispatchUncertain($siteId, $claim->previous['operation'], $claim->previous['batches'])) {
            $this->claims->announce($claim, $claim->previous['operation'], $claim->previous['batches']); // keep what is in doubt on the blocked task
            $this->claims->block($claim);
            return 'blocked';
        }
        $this->generations->discardProcessing($siteId, $generationId, $sourceId); // an earlier owner's unfinished version
        $size = $this->db->selectOne('SELECT COUNT(*) AS n, COALESCE(SUM(LENGTH(text)), 0) AS bytes FROM {knowledge_chunks} WHERE site_id = ? AND version_id = ?', [$siteId, (int) $current['id']]);
        return ['raw' => $text['raw'], 'type' => $text['type'], 'hash' => (string) $current['content_hash'], 'units' => (int) $size['bytes'] + self::PER_TEXT_OVERHEAD * (int) $size['n'], 'chunks' => (int) $size['n'], 'claim' => $claim];
    }

    /**
     * True when a request of this operation may have reached the provider without a known outcome (dispatched or uncertain in the
     * ledger). Requests that were only reserved (never sent) are released; dispatched ones are marked uncertain. Without a ledger the
     * answer for a named operation is "uncertain" only when the caller already knows a dispatch happened, so this returns false.
     */
    private function dispatchUncertain(int $siteId, string $operation, int $batches): bool
    {
        if ($batches < 1) {
            return false;
        }
        $ids = [];
        for ($n = 0; $n < $batches; $n++) {
            $ids[] = RemoteEmbedder::requestId(Ingestor::batchOperation($operation, $n));
        }
        $rows = $this->db->select('SELECT request_id, state FROM {reservations} WHERE site_id = ? AND request_id IN (' . implode(', ', array_fill(0, count($ids), '?')) . ')', array_merge([$siteId], $ids));
        $unsafe = false;
        foreach ($rows as $row) {
            switch ((string) $row['state']) {
                case 'reserved':
                    $this->ledger?->release($siteId, (string) $row['request_id'], 'worker-lost');
                    break;
                case 'dispatched':
                    $this->ledger?->markUncertain($siteId, (string) $row['request_id'], 'worker-lost');
                    $unsafe = true;
                    break;
                case 'uncertain':
                    $unsafe = true;
                    break;
            }
        }
        return $unsafe;
    }

    /**
     * Maps an ingestor rejection to pause or failure by its stable reason code (never by message).
     *
     * @return array{0:string, 1:string}
     */
    private function rejection(string $reason): array
    {
        return match ($reason) {
            'embedding_budget', 'embedding_failed', 'embedding_duplicate' => ['paused', $reason],
            'embedding_context', 'embedding_input' => ['failed', $reason],
            default => ['failed', 'source_rejected'],
        };
    }

    /**
     * Ends a step: stores the cursor, refreshes the progress, applies a pause/failure, or decides between continue, waiting and complete/synced.
     *
     * @param array<string, mixed> $out
     * @return array<string, mixed>
     */
    private function finishStep(int $siteId, int $generationId, bool $active, array $out, int $cursor, ?string $state, ?string $code): array
    {
        $this->generations->setCursor($siteId, $generationId, $cursor);
        $progress = $this->generations->refreshProgress($siteId, $generationId);
        $out['remaining'] = max(0, $progress['total'] - $progress['built']); // one count per step, reused
        if ($state === 'failed') {
            $r = $active ? $this->pause($siteId, $generationId, (string) $code) : $this->giveUp($siteId, $generationId, (string) $code);
            return ['state' => $r['state'], 'code' => $r['code']] + $out;
        }
        if ($state !== null) {
            if ($state === 'paused' || $state === 'retry') {
                $this->generations->note($siteId, $generationId, $code);
            }
            return ['state' => $state, 'code' => $code] + $out;
        }
        if ($out['remaining'] === 0) {
            if ($active) {
                $this->generations->note($siteId, $generationId, null);
                return ['state' => 'synced', 'code' => null] + $out;
            }
            try {
                $this->generations->note($siteId, $generationId, null);
                $this->generations->finish($siteId, $generationId);
            } catch (GenerationException $e) {
                // A source arrived, changed or restarted between the last check and now, or another worker is mid-source: go round again.
                if (in_array($e->getMessage(), ['incomplete', 'unfinished_versions'], true)) {
                    return ['state' => 'continue', 'code' => null] + $out;
                }
                if ($e->getMessage() === 'not_building') {
                    return ['state' => 'abandoned', 'code' => 'not_building'] + $out;
                }
                return $this->giveUp($siteId, $generationId, $e->getMessage()) + $out;
            }
            return ['state' => 'complete', 'code' => null] + $out;
        }
        if ($out['built'] + $out['skipped'] > 0) {
            return ['state' => 'continue', 'code' => null] + $out;
        }
        if ($out['blocked'] > 0 && $out['deferred'] === 0) {
            $this->generations->note($siteId, $generationId, 'uncertain_dispatch');
            return ['state' => 'paused', 'code' => 'uncertain_dispatch'] + $out;
        }
        return ['state' => 'waiting', 'code' => 'sources_not_settled'] + $out;
    }

    /** @return array<string, mixed> */
    private function giveUp(int $siteId, int $generationId, string $code): array
    {
        try {
            $this->generations->fail($siteId, $generationId, $code);
        } catch (GenerationException) {
            return $this->result('abandoned', 'not_building'); // activated or already ended meanwhile: nothing to give up
        }
        return $this->result('failed', $code);
    }

    /** @return array<string, mixed> */
    private function pause(int $siteId, int $generationId, string $code): array
    {
        $this->generations->note($siteId, $generationId, $code);
        return $this->result('paused', $code);
    }

    private function siteChunks(int $siteId): int
    {
        return (int) $this->db->selectOne(
            "SELECT COALESCE(SUM(cv.chunk_count), 0) AS n FROM {knowledge_sources} s JOIN {knowledge_versions} cv ON cv.id = s.current_version_id
             WHERE s.site_id = ? AND s.status = 'ready' AND s.visibility = 'public'",
            [$siteId]
        )['n'];
    }

    /** @return array{state:string, code:?string, built:int, skipped:int, deferred:int, blocked:int, remaining:int, units:int, source:?int, limit:?int, chunks:?int} */
    private function result(string $state, ?string $code): array
    {
        return ['state' => $state, 'code' => $code, 'built' => 0, 'skipped' => 0, 'deferred' => 0, 'blocked' => 0, 'remaining' => 0, 'units' => 0, 'source' => null, 'limit' => null, 'chunks' => null];
    }
}