package gudusoft.gsqlparser.pp2.island;

import gudusoft.gsqlparser.ETokenType;
import gudusoft.gsqlparser.pp2.island.SqlScopeDetector.SqlScopeResult;
import gudusoft.gsqlparser.pp2.token.Pp2Token;
import gudusoft.gsqlparser.pp2.token.Pp2TokenStream;
import gudusoft.gsqlparser.pp2.token.TokenRole;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;

/**
 * Annotates each token with the clause it belongs to ({@link ClausePart}) and
 * marks master/clause keyword roles, by walking the token stream left-to-right
 * and tracking the current clause per SQL scope level.
 *
 * <h2>SQL-scope aware</h2>
 *
 * <p>The annotator consumes {@link SqlScopeResult} (S20) so a subquery's
 * clauses do not leak into the enclosing query: the "current clause" is tracked
 * independently per scope level. When the walk descends into a deeper scope,
 * that level's clause state starts fresh; when it returns, the enclosing level's
 * clause (e.g. {@code FROM} around a derived-table subquery) is intact.
 *
 * <h2>Expression-depth gating</h2>
 *
 * <p>A clause keyword only changes the current clause when it appears at the
 * query block's own paren depth — not inside a nested expression. This stops
 * window specs ({@code OVER (ORDER BY a)}), aggregate filters
 * ({@code FILTER (WHERE ...)}), and parenthesised predicates from corrupting
 * the clause state. The function-form keywords {@code LEFT(} / {@code RIGHT(}
 * are also recognised as function calls (not JOIN keywords) when immediately
 * followed by an open paren.
 *
 * <h2>Designated annotator</h2>
 *
 * <p>S21 is one of the designated role-annotator stages (plan forbidden #6), so
 * it sets {@link TokenRole#KEYWORD_MASTER} on master keywords
 * (SELECT/INSERT/UPDATE/DELETE/MERGE) and {@link TokenRole#KEYWORD_CLAUSE} on
 * recognised clause keywords. It never mutates the wrapped {@code TSourceToken}.
 * The per-token clause membership is returned as a {@link ClauseScopeResult}
 * side structure (a token belongs to exactly one clause, so a small enum is
 * clearer than seven boolean roles).
 *
 * <p>Iterative single pass; no recursion. Plan reference: §7.3/S21, §7.4/S21.
 */
public final class ClauseScopeAnnotator {

    private static final Set<String> MASTER_KEYWORDS = setOf(
        "SELECT", "INSERT", "UPDATE", "DELETE", "MERGE");

    /** Keyword text -> the clause it starts. */
    private static final Map<String, ClausePart> CLAUSE_KEYWORDS = buildClauseKeywords();

    private static Map<String, ClausePart> buildClauseKeywords() {
        Map<String, ClausePart> m = new HashMap<String, ClausePart>();
        m.put("FROM", ClausePart.FROM);
        m.put("WHERE", ClausePart.WHERE);
        m.put("GROUP", ClausePart.GROUP_BY);
        m.put("ORDER", ClausePart.ORDER_BY);
        m.put("HAVING", ClausePart.HAVING);
        m.put("INTO", ClausePart.INTO);
        m.put("VALUES", ClausePart.INSERT_VALUES);
        m.put("SET", ClausePart.UPDATE_SET);
        // JOIN family — all map to the JOIN clause.
        for (String j : new String[]{"JOIN", "INNER", "LEFT", "RIGHT", "FULL",
            "CROSS", "NATURAL", "OUTER", "ON", "USING", "APPLY"}) {
            m.put(j, ClausePart.JOIN);
        }
        // Set operators.
        for (String s : new String[]{"UNION", "INTERSECT", "EXCEPT", "MINUS"}) {
            m.put(s, ClausePart.SET_OP);
        }
        return Collections.unmodifiableMap(m);
    }

    /**
     * Annotate {@code stream} clause membership using the SQL scope levels.
     *
     * @param stream   the token stream; must not be null
     * @param sqlScope the S20 scope result for the same stream; must not be null
     *                 and must cover the same token count
     * @return a {@link ClauseScopeResult}; never null
     * @throws NullPointerException if any argument is null
     * @throws IllegalArgumentException if {@code sqlScope.size()} != stream size
     */
    public ClauseScopeResult annotate(Pp2TokenStream stream, SqlScopeResult sqlScope) {
        if (stream == null) throw new NullPointerException("stream");
        if (sqlScope == null) throw new NullPointerException("sqlScope");
        int n = stream.size();
        if (sqlScope.size() != n) {
            throw new IllegalArgumentException(
                "sqlScope size " + sqlScope.size() + " != stream size " + n);
        }

        ClausePart[] part = new ClausePart[n];
        List<Integer> masterIndices = new ArrayList<Integer>();

        int levels = Math.max(1, sqlScope.getMaxLevel() + 1);
        // Current clause per SQL scope level. Index by level.
        ClausePart[] current = new ClausePart[levels];
        java.util.Arrays.fill(current, ClausePart.NONE);
        // The paren depth at which each level's clauses live (the query block's
        // base). Clause keywords only transition when paren depth == base[level].
        int[] base = new int[levels];

        int parenDepth = 0;
        int prevLevel = 0;
        for (int i = 0; i < n; i++) {
            Pp2Token t = stream.get(i);
            int level = sqlScope.levelAt(i);

            // Descending into deeper subquery scope(s) resets each newly entered
            // level's clause and records its base paren depth (just inside the
            // query paren). Reset every level in (prevLevel, level], not just the
            // target, in case S20 levels ever step by more than one.
            if (level > prevLevel) {
                for (int l = prevLevel + 1; l <= level; l++) {
                    current[l] = ClausePart.NONE;
                    base[l] = parenDepth;
                }
            }
            prevLevel = level;

            ETokenType type = t.getSourceToken().tokentype;
            if (type == ETokenType.ttsemicolon) {
                // Statement boundary: reset clause state so the next statement's
                // header does not inherit the previous statement's clause.
                java.util.Arrays.fill(current, ClausePart.NONE);
                java.util.Arrays.fill(base, 0);
                parenDepth = 0;
                part[i] = ClausePart.NONE;
            } else if (type == ETokenType.ttleftparenthesis) {
                part[i] = current[level];
                parenDepth++;
            } else if (type == ETokenType.ttrightparenthesis) {
                if (parenDepth > 0) parenDepth--;
                part[i] = current[level];
            } else if (type == ETokenType.ttkeyword) {
                String text = t.getText();
                String upper = text == null ? "" : text.toUpperCase(Locale.ROOT);
                boolean atBase = parenDepth == base[level];
                if (MASTER_KEYWORDS.contains(upper)) {
                    t.addRole(TokenRole.KEYWORD_MASTER);
                    masterIndices.add(i);
                    if (atBase) {
                        // SELECT opens its projection list; other masters leave the
                        // clause NONE until a recognised sub-clause keyword appears.
                        current[level] = "SELECT".equals(upper)
                            ? ClausePart.SELECT_LIST : ClausePart.NONE;
                    }
                    part[i] = current[level];
                } else if (CLAUSE_KEYWORDS.containsKey(upper)
                        && !isFunctionForm(upper, stream, i)
                        && transitionAllowed(CLAUSE_KEYWORDS.get(upper), current[level],
                                             atBase)) {
                    // A genuine clause transition.
                    t.addRole(TokenRole.KEYWORD_CLAUSE);
                    current[level] = CLAUSE_KEYWORDS.get(upper);
                    part[i] = current[level];
                } else {
                    // Generic keyword, OR a clause keyword inside an expression
                    // (window/filter/func/parenthesised predicate): stays in the
                    // current clause without transitioning.
                    part[i] = current[level];
                }
            } else {
                part[i] = current[level];
            }
        }

        return new ClauseScopeResult(part, Collections.unmodifiableList(masterIndices));
    }

    /**
     * Whether a clause keyword may transition the current clause given the
     * paren context. Transitions normally require being at the query block's
     * base paren depth ({@code atBase}). The exception is JOIN-family keywords
     * inside a table-expression group such as {@code FROM (t1 LEFT JOIN t2 ON ...)}:
     * those legitimately appear above base, but only when the current clause is
     * already FROM or JOIN (the table region). Window/filter/expression keywords
     * are not JOIN-family, so they remain gated by base depth.
     */
    private static boolean transitionAllowed(ClausePart mapped, ClausePart current,
                                             boolean atBase) {
        if (atBase) return true;
        return mapped == ClausePart.JOIN
            && (current == ClausePart.FROM || current == ClausePart.JOIN);
    }

    /** Keywords that double as function names; treated as functions when followed by '('. */
    private static final Set<String> FUNCTION_FORM_KEYWORDS = setOf("LEFT", "RIGHT");

    /**
     * True if {@code upper} is a keyword that doubles as a function name (LEFT,
     * RIGHT) and is immediately followed (skipping comments) by an open paren —
     * i.e. it is a function call, not a JOIN keyword. Other clause keywords that
     * legitimately precede a paren (FROM/JOIN/WHERE/ON/APPLY before a subquery or
     * parenthesised predicate) are not treated as function forms.
     */
    private static boolean isFunctionForm(String upper, Pp2TokenStream stream, int i) {
        if (!FUNCTION_FORM_KEYWORDS.contains(upper)) return false;
        for (int j = i + 1; j < stream.size(); j++) {
            ETokenType type = stream.get(j).getSourceToken().tokentype;
            if (type == ETokenType.ttsimplecomment
                || type == ETokenType.ttbracketedcomment
                || type == ETokenType.ttCPPComment) {
                continue;
            }
            return type == ETokenType.ttleftparenthesis;
        }
        return false;
    }

    private static Set<String> setOf(String... values) {
        Set<String> s = new HashSet<String>();
        Collections.addAll(s, values);
        return Collections.unmodifiableSet(s);
    }

    /** Per-token clause membership plus the master-keyword token indices. */
    public static final class ClauseScopeResult {
        private final ClausePart[] part;
        private final List<Integer> masterIndices;

        ClauseScopeResult(ClausePart[] part, List<Integer> masterIndices) {
            this.part = part;
            this.masterIndices = masterIndices;
        }

        /** Number of tokens covered. */
        public int size() { return part.length; }

        /** The clause the token at {@code index} belongs to. */
        public ClausePart partAt(int index) { return part[index]; }

        /** Indices of master keyword tokens (SELECT/INSERT/UPDATE/DELETE/MERGE), in order. */
        public List<Integer> getMasterIndices() { return masterIndices; }
    }
}