/*
 *  Licensed to the Apache Software Foundation (ASF) under one
 *  or more contributor license agreements.  See the NOTICE file
 *  distributed with this work for additional information
 *  regarding copyright ownership.  The ASF licenses this file
 *  to you under the Apache License, Version 2.0 (the
 *  "License"); you may not use this file except in compliance
 *  with the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing,
 *  software distributed under the License is distributed on an
 *  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 *  KIND, either express or implied.  See the License for the
 *  specific language governing permissions and limitations
 *  under the License.
 */
package org.apache.groovy.contracts.ast;

import groovy.contracts.Decreases;
import groovy.transform.CompilationUnitAware;
import org.apache.groovy.contracts.VariantSupport;
import org.codehaus.groovy.ast.ASTNode;
import org.codehaus.groovy.ast.AnnotationNode;
import org.codehaus.groovy.ast.ClassHelper;
import org.codehaus.groovy.ast.expr.ClosureExpression;
import org.codehaus.groovy.ast.expr.Expression;
import org.codehaus.groovy.ast.stmt.BlockStatement;
import org.codehaus.groovy.ast.stmt.ExpressionStatement;
import org.codehaus.groovy.ast.stmt.LoopingStatement;
import org.codehaus.groovy.ast.stmt.Statement;
import org.codehaus.groovy.control.CompilationUnit;
import org.codehaus.groovy.control.CompilePhase;
import org.codehaus.groovy.control.SourceUnit;
import org.codehaus.groovy.transform.ASTTransformation;
import org.codehaus.groovy.transform.GroovyASTTransformation;

import java.util.List;
import java.util.concurrent.atomic.AtomicLong;

import static org.codehaus.groovy.ast.tools.GeneralUtils.args;
import static org.codehaus.groovy.ast.tools.GeneralUtils.assignS;
import static org.codehaus.groovy.ast.tools.GeneralUtils.block;
import static org.codehaus.groovy.ast.tools.GeneralUtils.boolX;
import static org.codehaus.groovy.ast.tools.GeneralUtils.callX;
import static org.codehaus.groovy.ast.tools.GeneralUtils.constX;
import static org.codehaus.groovy.ast.tools.GeneralUtils.declS;
import static org.codehaus.groovy.ast.tools.GeneralUtils.ifElseS;
import static org.codehaus.groovy.ast.tools.GeneralUtils.localVarX;
import static org.codehaus.groovy.ast.tools.GeneralUtils.nullX;
import static org.codehaus.groovy.ast.tools.GeneralUtils.stmt;
import static org.codehaus.groovy.ast.tools.GeneralUtils.varX;

/**
 * Handles {@link Decreases} annotations placed on loop statements ({@code for},
 * {@code while}, {@code do-while}). The closure must return a value that is
 * non-negative at the start of each iteration and strictly decreases between
 * consecutive iterations.
 * <p>
 * The transformation injects code to:
 * <ol>
 *   <li>Evaluate the expression at the start of each iteration.</li>
 *   <li>On the first iteration, assert a scalar measure is non-negative at
 *       loop entry (well-foundedness).</li>
 *   <li>On subsequent iterations, assert it has strictly decreased since the
 *       previous iteration start, and is non-negative (GROOVY-12128: comparing
 *       iteration starts — rather than the start and end of one pass — sees
 *       progress made by a classic {@code for} loop's update expression, which
 *       runs after the body, and does not demand non-negativity after the
 *       final body execution).</li>
 * </ol>
 * The final iteration's decrease is intentionally unchecked: a loop that has
 * exited needs no further termination evidence.
 * <p>
 * Example:
 * <pre>
 * int n = 10
 * {@code @Decreases}({ n })
 * while (n &gt; 0) {
 *     n--
 * }
 * </pre>
 *
 * @since 6.0.0
 * @see Decreases
 * @see org.apache.groovy.contracts.VariantSupport
 * @see org.apache.groovy.contracts.LoopVariantViolation
 */
@GroovyASTTransformation(phase = CompilePhase.SEMANTIC_ANALYSIS)
public class LoopVariantASTTransformation implements ASTTransformation, CompilationUnitAware {

    private static final AtomicLong COUNTER = new AtomicLong();

    private CompilationUnit compilationUnit;

    @Override
    public void setCompilationUnit(final CompilationUnit unit) {
        this.compilationUnit = unit;
    }

    /**
     * Rewrites a loop-level {@link Decreases} annotation into variant bookkeeping and runtime checks.
     *
     * @param nodes the annotated AST nodes supplied by the compiler
     * @param source the current source unit
     */
    @Override
    public void visit(final ASTNode[] nodes, final SourceUnit source) {
        if (nodes.length != 2) return;
        if (!(nodes[0] instanceof AnnotationNode annotation)) return;
        if (!(nodes[1] instanceof LoopingStatement loopStatement)) return;

        Expression value = annotation.getMember("value");
        if (!(value instanceof ClosureExpression closureExpression)) return;

        Expression variantExpression = extractExpression(closureExpression);
        if (variantExpression == null) return;

        String suffix = Long.toString(COUNTER.getAndIncrement());
        String prevVarName = "$_gc_decreases_prev_" + suffix;
        String currVarName = "$_gc_decreases_curr_" + suffix;
        String seenVarName = "$_gc_decreases_seen_" + suffix;

        String className = LoopContractSupport.enclosingClassName(source, (ASTNode) loopStatement);
        BlockStatement enclosingBlock = LoopContractSupport.enclosingBlock(source, (ASTNode) loopStatement);

        if (enclosingBlock != null) {
            // GROOVY-12128: measure the variant at the start of each iteration and compare
            // consecutive iteration-start values. Measuring within a single pass (body start vs
            // body end) misses progress made by a classic for loop's update expression, which
            // runs after the body end, and wrongly demands non-negativity after the final body
            // execution instead of at iteration entry.
            Statement declPrev = declS(localVarX(prevVarName, ClassHelper.dynamicType()), nullX());
            declPrev.setSourcePosition(annotation);
            Statement declSeen = declS(localVarX(seenVarName, ClassHelper.boolean_TYPE), constX(false, true));
            declSeen.setSourcePosition(annotation);
            List<Statement> siblings = enclosingBlock.getStatements();
            siblings.addAll(siblings.indexOf((Statement) loopStatement), List.of(declPrev, declSeen));

            Statement saveCurr = declS(localVarX(currVarName, ClassHelper.dynamicType()), variantExpression);
            saveCurr.setSourcePosition(annotation);
            // on the first iteration there is no previous value to compare against, but a scalar
            // measure must already be non-negative at loop entry (well-foundedness)
            Statement decreaseCheck = ifElseS(
                    boolX(varX(seenVarName)),
                    stmt(callX(
                            ClassHelper.make(VariantSupport.class),
                            "checkLoopVariant",
                            args(varX(prevVarName), varX(currVarName), constX(className))
                    )),
                    stmt(callX(
                            ClassHelper.make(VariantSupport.class),
                            "checkLoopVariantEntry",
                            args(varX(currVarName), constX(className))
                    ))
            );
            decreaseCheck.setSourcePosition(annotation);
            Statement markSeen = assignS(varX(seenVarName), constX(true, true));
            markSeen.setSourcePosition(annotation);
            Statement savePrev = assignS(varX(prevVarName), varX(currVarName));
            savePrev.setSourcePosition(annotation);

            injectAtLoopBodyStart(loopStatement, block(saveCurr, decreaseCheck, markSeen, savePrev));
        } else {
            // The loop is not a direct child of a block (e.g. a braceless if branch), so there is
            // nowhere to hoist cross-iteration state; fall back to the single-pass measurement.

            // At start of iteration: def prevVar = <expression>
            Statement savePrev = declS(localVarX(prevVarName, ClassHelper.dynamicType()), variantExpression);
            savePrev.setSourcePosition(annotation);

            // At end of iteration: def currVar = <expression copy>
            // We need a fresh copy of the expression for re-evaluation
            Expression variantCopy = copyExpression(closureExpression);
            Statement saveCurr = declS(localVarX(currVarName, ClassHelper.dynamicType()), variantCopy);
            saveCurr.setSourcePosition(annotation);

            // Assert (if enabled): currVar < prevVar and non-negative — delegated to the
            // shared VariantSupport, gated by the enclosing class's -ea/-da configuration.
            Statement decreaseCheck = stmt(
                    callX(
                            ClassHelper.make(VariantSupport.class),
                            "checkLoopVariant",
                            args(varX(prevVarName), varX(currVarName), constX(className))
                    )
            );
            decreaseCheck.setSourcePosition(annotation);

            // Inject: save at start, check at end
            injectAtLoopBodyStartAndEnd(loopStatement, savePrev, block(saveCurr, decreaseCheck));
        }

        // The variant closure lived inside a statement annotation, so the compiler's resolution
        // passes never reached it; re-resolve types, static imports and variable scopes now that the
        // expressions are real loop-body statements.
        LoopContractSupport.resolveInlinedContractCode(source, (ASTNode) loopStatement, compilationUnit);
    }

    private static Expression extractExpression(ClosureExpression closureExpression) {
        BlockStatement block = (BlockStatement) closureExpression.getCode();
        List<Statement> statements = block.getStatements();
        if (statements.size() != 1) return null;
        Statement stmt = statements.get(0);
        if (stmt instanceof ExpressionStatement) {
            return ((ExpressionStatement) stmt).getExpression();
        }
        return null;
    }

    private static Expression copyExpression(ClosureExpression closureExpression) {
        // Re-extract from the closure to get a fresh AST node
        // (the original is consumed by the first injection point)
        BlockStatement block = (BlockStatement) closureExpression.getCode();
        List<Statement> statements = block.getStatements();
        if (statements.size() != 1) return null;
        Statement stmt = statements.get(0);
        if (stmt instanceof ExpressionStatement exprStmt) {
            // Use transformExpression to get a deep copy
            return exprStmt.getExpression().transformExpression(expr -> expr);
        }
        return null;
    }

    private static void injectAtLoopBodyStart(LoopingStatement loopStatement, Statement bookkeeping) {
        Statement loopBody = loopStatement.getLoopBlock();
        if (loopBody instanceof BlockStatement block) {
            block.getStatements().addAll(0, ((BlockStatement) bookkeeping).getStatements());
        } else {
            BlockStatement newBody = new BlockStatement();
            newBody.addStatements(((BlockStatement) bookkeeping).getStatements());
            newBody.addStatement(loopBody);
            newBody.setSourcePosition(loopBody);
            loopStatement.setLoopBlock(newBody);
        }
    }

    private static void injectAtLoopBodyStartAndEnd(LoopingStatement loopStatement,
                                                     Statement startCheck, Statement endCheck) {
        Statement loopBody = loopStatement.getLoopBlock();
        BlockStatement newBody;
        if (loopBody instanceof BlockStatement block) {
            // Prepend save at start
            block.getStatements().add(0, startCheck);
            // Append checks at end
            block.getStatements().addAll(((BlockStatement) endCheck).getStatements());
            newBody = block;
        } else {
            newBody = new BlockStatement();
            newBody.addStatement(startCheck);
            newBody.addStatement(loopBody);
            newBody.addStatements(((BlockStatement) endCheck).getStatements());
            newBody.setSourcePosition(loopBody);
            loopStatement.setLoopBlock(newBody);
        }
    }
}
