<?php

declare(strict_types=1);

namespace Antlr\Antlr4\Runtime\Atn;

use Antlr\Antlr4\Runtime\Atn\Actions\LexerAction;
use Antlr\Antlr4\Runtime\Atn\States\ATNState;
use Antlr\Antlr4\Runtime\Atn\States\DecisionState;
use Antlr\Antlr4\Runtime\Atn\States\RuleStartState;
use Antlr\Antlr4\Runtime\Atn\States\RuleStopState;
use Antlr\Antlr4\Runtime\Atn\States\TokensStartState;
use Antlr\Antlr4\Runtime\Atn\Transitions\RuleTransition;
use Antlr\Antlr4\Runtime\IntervalSet;
use Antlr\Antlr4\Runtime\LL1Analyzer;
use Antlr\Antlr4\Runtime\ParserRuleContext;
use Antlr\Antlr4\Runtime\RuleContext;
use Antlr\Antlr4\Runtime\Token;

final class ATN
{
    public const INVALID_ALT_NUMBER = 0;
    public const ATN_TYPE_LEXER = 0;
    public const ATN_TYPE_PARSER = 1;

    /** @var array<ATNState> */
    public array $states = [];

    /**
     * Each subrule/rule is a decision point and we must track them so we
     * can go back later and build DFA predictors for them. This includes
     * all the rules, subrules, optional blocks, ()+, ()* etc...
     *
     * @var array<DecisionState>
     */
    public array $decisionToState = [];

    /**
     * Maps from rule index to starting state number.
     *
     * @var array<RuleStartState>
     */
    public array $ruleToStartState = [];

    /**
     * Maps from rule index to stop state number.
     *
     * @var array<RuleStopState>
     */
    public array $ruleToStopState = [];

    /** @var array<TokensStartState> */
    public array $modeNameToStartState = [];

    /**
     * The type of the ATN.
     */
    public int $grammarType;

    /**
     * The maximum value for any symbol recognized by a transition in the ATN.
     */
    public int $maxTokenType;

    /**
     * For lexer ATNs, this maps the rule index to the resulting token type.
     * For parser ATNs, this maps the rule index to the generated bypass token
     * type if the
     * {@see ATNDeserializationOptions::isGenerateRuleBypassTransitions()}
     * deserialization option was specified; otherwise, this is `null`.
     *
     * @var array<int|null>
     */
    public array $ruleToTokenType = [];

    /**
     * For lexer ATNs, this is an array of {@see LexerAction} objects which may
     * be referenced by action transitions in the ATN.
     *
     * @var array<LexerAction>
     */
    public array $lexerActions = [];

    /** @var array<TokensStartState> */
    public array $modeToStartState = [];

    /**
     * Used for runtime deserialization of ATNs from strings
     */
    public function __construct(int $grammarType, int $maxTokenType)
    {
        $this->grammarType = $grammarType;
        $this->maxTokenType = $maxTokenType;
    }

    /**
     * Compute the set of valid tokens that can occur starting in state `state`.
     * If `context` is null, the set of tokens will not include what can follow
     * the rule surrounding `state`. In other words, the set will be
     * restricted to tokens reachable staying within `state`'s rule.
     */
    public function nextTokensInContext(ATNState $state, ?RuleContext $context): IntervalSet
    {
        return (new LL1Analyzer($this))->look($state, null, $context);
    }

    /**
     * Compute the set of valid tokens that can occur starting in `state` and
     * staying in same rule. {@see Token::EPSILON} is in set if we reach end of
     * rule.
     */
    public function nextTokens(ATNState $state): IntervalSet
    {
        if ($state->nextTokenWithinRule !== null) {
            return $state->nextTokenWithinRule;
        }

        $state->nextTokenWithinRule = $this->nextTokensInContext($state, null);
        $state->nextTokenWithinRule->setReadOnly(true);

        return $state->nextTokenWithinRule;
    }

    public function addState(?ATNState $state): void
    {
        if ($state === null) {
            return;
        }

        $state->atn = $this;
        $state->stateNumber = \count($this->states);

        $this->states[] = $state;
    }

    public function removeState(ATNState $state): void
    {
        // just free mem, don't shift states in list
        unset($this->states[$state->stateNumber]);
    }

    public function defineDecisionState(DecisionState $s): int
    {
        $this->decisionToState[] = $s;

        $s->decision = \count($this->decisionToState) - 1;

        return $s->decision;
    }

    public function getDecisionState(int $decision): ?DecisionState
    {
        if (\count($this->decisionToState) === 0) {
            return null;
        }

        return $this->decisionToState[$decision];
    }

    public function getNumberOfDecisions(): int
    {
        return \count($this->decisionToState);
    }

    /**
     * Computes the set of input symbols which could follow ATN state number
     * `stateNumber` in the specified full `context`. This method
     * considers the complete parser context, but does not evaluate semantic
     * predicates (i.e. all predicates encountered during the calculation are
     * assumed true). If a path in the ATN exists from the starting state to the
     * {@see RuleStopState} of the outermost context without matching any
     * symbols, {@see Token::EOF} is added to the returned set.
     *
     * If `context` is `null`, it is treated as {@see ParserRuleContext::EMPTY}.
     *
     * @param int         $stateNumber The ATN state number
     * @param RuleContext $context     The full parse context
     *
     * @return IntervalSet The set of potentially valid input symbols which could
     *                     follow the specified state in the specified context.
     *
     * @throws \LogicException
     */
    public function getExpectedTokens(int $stateNumber, ?RuleContext $context): IntervalSet
    {
        if ($stateNumber < 0 || $stateNumber >= \count($this->states)) {
            throw new \InvalidArgumentException('Invalid state number.');
        }

        $s = $this->states[$stateNumber];
        $following = $this->nextTokens($s);

        if (!$following->contains(Token::EPSILON)) {
            return $following;
        }

        $expected = new IntervalSet();

        $expected->addSet($following);
        $expected->removeOne(Token::EPSILON);

        if ($context === null) {
            $context = ParserRuleContext::emptyContext();
        }

        while ($context !== null && $context->invokingState >= 0 && $following->contains(Token::EPSILON)) {
            $invokingState = $this->states[$context->invokingState];
            $transition = $invokingState->getTransition(0);

            if (!$transition instanceof RuleTransition) {
                throw new \LogicException('Unexpected transition type.');
            }

            $following = $this->nextTokens($transition->followState);

            $expected->addSet($following);
            $expected->removeOne(Token::EPSILON);

            $context = $context->getParent();
        }

        if ($following->contains(Token::EPSILON)) {
            $expected->addOne(Token::EOF);
        }

        return $expected;
    }
}
