package click.kamil.springstatemachineexporter.analysis.pipeline;

import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
import org.eclipse.jdt.core.dom.ASTVisitor;
import org.eclipse.jdt.core.dom.Expression;
import org.eclipse.jdt.core.dom.MethodDeclaration;
import org.eclipse.jdt.core.dom.MethodInvocation;
import org.eclipse.jdt.core.dom.ParenthesizedExpression;
import org.eclipse.jdt.core.dom.ReturnStatement;
import org.eclipse.jdt.core.dom.SingleVariableDeclaration;
import org.eclipse.jdt.core.dom.TypeDeclaration;

import java.util.ArrayList;
import java.util.List;
import java.util.function.Function;

/**
 * Shared {@link MethodInvocation} peeling for factory/reactive wrappers and local passthrough methods.
 */
public final class MethodInvocationUnwrapper {

    private MethodInvocationUnwrapper() {
    }

    public static Expression unwrap(
            MethodInvocation mi,
            int depth,
            String methodFqn,
            CodebaseContext context,
            Function<MethodInvocation, String> targetMethodFqnResolver) {
        if (ResolutionBudget.defaults().isUnwrapExhausted(depth)) {
            return mi;
        }

        String targetMethodFqn = targetMethodFqnResolver != null
                ? targetMethodFqnResolver.apply(mi)
                : null;
        if (targetMethodFqn == null && methodFqn != null && methodFqn.contains(".") && mi.getExpression() == null) {
            String className = methodFqn.substring(0, methodFqn.lastIndexOf('.'));
            targetMethodFqn = className + "." + mi.getName().getIdentifier();
        }

        if (targetMethodFqn != null && targetMethodFqn.contains(".")) {
            String className = targetMethodFqn.substring(0, targetMethodFqn.lastIndexOf('.'));
            String methodName = targetMethodFqn.substring(targetMethodFqn.lastIndexOf('.') + 1);

            TypeDeclaration currentTd = methodFqn != null && methodFqn.contains(".")
                    ? context.getTypeDeclaration(methodFqn.substring(0, methodFqn.lastIndexOf('.')))
                    : null;
            org.eclipse.jdt.core.dom.CompilationUnit cu = currentTd != null && currentTd.getRoot() instanceof org.eclipse.jdt.core.dom.CompilationUnit compilationUnit
                    ? compilationUnit
                    : null;

            TypeDeclaration td = context.getTypeDeclaration(className, cu);
            if (td == null) {
                td = context.getTypeDeclaration(className);
            }
            if (td != null) {
                MethodDeclaration localMd = context.findMethodDeclaration(td, methodName, true);
                if (localMd != null) {
                    int paramIdx = getReturnedParameterIndex(localMd);
                    if (paramIdx >= 0 && paramIdx < mi.arguments().size()) {
                        Expression arg = peelExpression((Expression) mi.arguments().get(paramIdx));
                        if (arg instanceof MethodInvocation innerMi) {
                            return unwrap(innerMi, depth + 1, methodFqn, context, targetMethodFqnResolver);
                        }
                        return arg;
                    }
                }
            }
        }

        if (mi.getExpression() != null) {
            String exprStr = mi.getExpression().toString();
            if (!LibraryUnwrapRegistry.isUnwrapReceiverType(exprStr)) {
                return mi;
            }
        } else if (targetMethodFqn == null) {
            return mi;
        }

        if (!mi.arguments().isEmpty()) {
            String methodName = mi.getName().getIdentifier();
            if (LibraryUnwrapRegistry.shouldStopUnwrap(methodName)) {
                return mi;
            }

            Expression arg = peelExpression((Expression) mi.arguments().get(0));
            if (arg instanceof MethodInvocation innerMi) {
                return unwrap(innerMi, depth + 1, methodFqn, context, targetMethodFqnResolver);
            }
            return arg;
        }
        return mi;
    }

    public static Expression peelExpression(Expression expression) {
        Expression current = expression;
        while (current instanceof ParenthesizedExpression pe) {
            current = pe.getExpression();
        }
        return current;
    }

    public static int getReturnedParameterIndex(MethodDeclaration md) {
        if (md.getBody() == null) {
            return -1;
        }
        List<?> parameters = md.parameters();
        if (parameters.isEmpty()) {
            return -1;
        }

        final List<String> returnedExprs = new ArrayList<>();
        md.getBody().accept(new ASTVisitor() {
            @Override
            public boolean visit(ReturnStatement node) {
                if (node.getExpression() != null) {
                    returnedExprs.add(node.getExpression().toString());
                }
                return super.visit(node);
            }
        });

        if (returnedExprs.size() == 1) {
            String retStr = returnedExprs.get(0);
            for (int i = 0; i < parameters.size(); i++) {
                if (parameters.get(i) instanceof SingleVariableDeclaration svd
                        && svd.getName().getIdentifier().equals(retStr)) {
                    return i;
                }
            }
        }
        return -1;
    }
}
