package click.kamil.springstatemachineexporter.analysis.pipeline;

import click.kamil.springstatemachineexporter.analysis.resolver.ConstantResolver;
import click.kamil.springstatemachineexporter.ast.common.AstUtils;
import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
import org.eclipse.jdt.core.dom.Block;
import org.eclipse.jdt.core.dom.Expression;
import org.eclipse.jdt.core.dom.LambdaExpression;
import org.eclipse.jdt.core.dom.MethodInvocation;
import org.eclipse.jdt.core.dom.ReturnStatement;
import org.eclipse.jdt.core.dom.SimpleName;
import org.eclipse.jdt.core.dom.ASTNode;

/**
 * Extracts terminal payload literals from reactive factory chains ({@code Mono.just}, nested transforms).
 */
public final class ReactiveExpressionSupport {

    private ReactiveExpressionSupport() {
    }

    public static String extractPayload(Expression expression, ConstantResolver constantResolver, CodebaseContext context) {
        return extractPayload(expression, null, constantResolver, context);
    }

    /**
     * Peels {@code just}/{@code withPayload}/{@code success} wrappers on an invocation chain and returns
     * the payload {@link Expression} for dataflow analysis.
     */
    public static Expression peelFactoryPayloadExpression(Expression expression) {
        if (!(expression instanceof MethodInvocation mi)) {
            return null;
        }
        MethodInvocation current = mi;
        while (current != null) {
            String methodName = current.getName().getIdentifier();
            if (LibraryUnwrapRegistry.isReactiveFactoryMethod(methodName) && !current.arguments().isEmpty()) {
                return MethodInvocationUnwrapper.peelExpression((Expression) current.arguments().get(0));
            }
            Expression receiver = current.getExpression();
            if (receiver instanceof MethodInvocation nextMi) {
                current = nextMi;
            } else {
                break;
            }
        }
        return null;
    }

    public static Expression lambdaBodyExpression(LambdaExpression lambda) {
        if (lambda.getBody() instanceof Expression bodyExpression) {
            return bodyExpression;
        }
        if (lambda.getBody() instanceof Block block) {
            for (Object statement : block.statements()) {
                if (statement instanceof ReturnStatement returnStatement
                        && returnStatement.getExpression() != null) {
                    return returnStatement.getExpression();
                }
            }
        }
        return null;
    }

    /**
     * Rewrites {@code p.getEvent()} inside a reactive transform lambda to {@code payload.getEvent()}
     * when {@code p} is fed by {@code Mono.just(payload)} (or a chained reactive source).
     */
    public static String remapLambdaParameterGetter(
            Expression eventExpr,
            ConstantResolver constantResolver,
            CodebaseContext context) {
        if (!(eventExpr instanceof MethodInvocation getterCall)) {
            return null;
        }
        if (!(getterCall.getExpression() instanceof SimpleName paramName)) {
            return null;
        }
        LambdaExpression lambda = findEnclosingLambda(getterCall);
        if (lambda == null) {
            return null;
        }
        MethodInvocation transformCall = findEnclosingReactiveTransform(lambda);
        if (transformCall == null) {
            return null;
        }
        String mappedReceiver = mapLambdaParameterToSource(
                lambda,
                paramName.getIdentifier(),
                transformCall.getExpression(),
                constantResolver,
                context);
        if (mappedReceiver == null) {
            return null;
        }
        return mappedReceiver + "." + getterCall.getName().getIdentifier() + "()";
    }

    /**
     * Same as {@link #remapLambdaParameterGetter(Expression, ConstantResolver, CodebaseContext)} but locates
     * the getter in a real method body when {@code expressionText} came from a detached parse tree.
     */
    public static String remapLambdaParameterGetterInMethod(
            String expressionText,
            String methodFqn,
            ConstantResolver constantResolver,
            CodebaseContext context) {
        MethodInvocation getter = AstUtils.findMethodInvocationInMethod(methodFqn, expressionText, context);
        if (getter == null) {
            return null;
        }
        String remapped = remapLambdaParameterGetter(getter, constantResolver, context);
        if (remapped != null) {
            return remapped;
        }
        return null;
    }

    private static LambdaExpression findEnclosingLambda(ASTNode node) {
        ASTNode current = node.getParent();
        while (current != null) {
            if (current instanceof LambdaExpression lambda) {
                return lambda;
            }
            current = current.getParent();
        }
        return null;
    }

    private static MethodInvocation findEnclosingReactiveTransform(ASTNode node) {
        ASTNode current = node.getParent();
        while (current != null) {
            if (current instanceof MethodInvocation mi
                    && LibraryUnwrapRegistry.isReactiveTransformMethod(mi.getName().getIdentifier())) {
                return mi;
            }
            current = current.getParent();
        }
        return null;
    }

    private static String extractPayload(
            Expression expression,
            Expression flatMapReceiver,
            ConstantResolver constantResolver,
            CodebaseContext context) {
        if (expression == null) {
            return null;
        }
        if (expression instanceof MethodInvocation mi) {
            String methodName = mi.getName().getIdentifier();
            if (LibraryUnwrapRegistry.isReactiveTransformMethod(methodName) && !mi.arguments().isEmpty()) {
                String fromArgument = extractTransformArgumentPayload(
                        (Expression) mi.arguments().get(0), mi.getExpression(), constantResolver, context);
                if (fromArgument != null) {
                    return fromArgument;
                }
            }
            if (LibraryUnwrapRegistry.isReactiveFactoryMethod(methodName) && !mi.arguments().isEmpty()) {
                Expression arg = (Expression) mi.arguments().get(0);
                if (constantResolver != null && context != null) {
                    String resolved = constantResolver.resolve(arg, context);
                    if (resolved != null) {
                        return resolved;
                    }
                }
                return arg.toString();
            }
            if (mi.getExpression() != null) {
                return extractPayload(mi.getExpression(), flatMapReceiver, constantResolver, context);
            }
        }
        return null;
    }

    private static String extractTransformArgumentPayload(
            Expression argument,
            Expression transformReceiver,
            ConstantResolver constantResolver,
            CodebaseContext context) {
        if (argument instanceof LambdaExpression lambda) {
            Expression bodyExpression = lambdaBodyExpression(lambda);
            if (bodyExpression != null) {
                String lambdaPayload = extractLambdaBodyPayload(
                        bodyExpression, lambda, transformReceiver, constantResolver, context);
                if (lambdaPayload != null) {
                    return lambdaPayload;
                }
                return extractPayload(bodyExpression, transformReceiver, constantResolver, context);
            }
        }
        return extractPayload(argument, transformReceiver, constantResolver, context);
    }

    private static String extractLambdaBodyPayload(
            Expression bodyExpression,
            LambdaExpression lambda,
            Expression transformReceiver,
            ConstantResolver constantResolver,
            CodebaseContext context) {
        if (!(bodyExpression instanceof MethodInvocation factoryCall)) {
            return null;
        }
        if (!LibraryUnwrapRegistry.isReactiveFactoryMethod(factoryCall.getName().getIdentifier())
                || factoryCall.arguments().isEmpty()) {
            return null;
        }
        Expression factoryArg = (Expression) factoryCall.arguments().get(0);
        if (!(factoryArg instanceof MethodInvocation getterCall)) {
            return null;
        }
        String getterSuffix = "." + getterCall.getName().getIdentifier() + "()";
        Expression getterReceiver = getterCall.getExpression();
        if (getterReceiver instanceof SimpleName paramName) {
            String mappedReceiver = mapLambdaParameterToSource(
                    lambda, paramName.getIdentifier(), transformReceiver, constantResolver, context);
            if (mappedReceiver != null) {
                return mappedReceiver + getterSuffix;
            }
        }
        if (getterReceiver != null) {
            return getterReceiver.toString() + getterSuffix;
        }
        return null;
    }

    private static String mapLambdaParameterToSource(
            LambdaExpression lambda,
            String paramName,
            Expression transformReceiver,
            ConstantResolver constantResolver,
            CodebaseContext context) {
        if (lambda.parameters().size() != 1) {
            return null;
        }
        Object parameter = lambda.parameters().get(0);
        if (!(parameter instanceof org.eclipse.jdt.core.dom.VariableDeclaration variableDeclaration)) {
            return null;
        }
        if (!paramName.equals(variableDeclaration.getName().getIdentifier())) {
            return null;
        }
        if (transformReceiver instanceof MethodInvocation receiverTransform
                && LibraryUnwrapRegistry.isReactiveTransformMethod(receiverTransform.getName().getIdentifier())
                && !receiverTransform.arguments().isEmpty()
                && receiverTransform.arguments().get(0) instanceof LambdaExpression feederLambda
                && feederLambda != lambda) {
            Expression feederBody = lambdaBodyExpression(feederLambda);
            if (feederBody != null) {
                String feederPayload = extractLambdaBodyPayload(
                        feederBody, feederLambda, receiverTransform.getExpression(), constantResolver, context);
                if (feederPayload != null) {
                    return feederPayload;
                }
                String extracted = extractPayload(
                        feederBody, receiverTransform.getExpression(), constantResolver, context);
                if (extracted != null) {
                    return extracted;
                }
            }
        }
        Expression source = peelJustArgument(transformReceiver);
        if (source instanceof MethodInvocation getterMi
                && getterMi.getExpression() instanceof SimpleName innerParamName) {
            LambdaExpression outerLambda = findEnclosingLambda(lambda);
            if (outerLambda != null && outerLambda != lambda) {
                MethodInvocation outerTransform = findEnclosingReactiveTransform(outerLambda);
                if (outerTransform != null) {
                    String mappedBase = mapLambdaParameterToSource(
                            outerLambda,
                            innerParamName.getIdentifier(),
                            outerTransform.getExpression(),
                            constantResolver,
                            context);
                    if (mappedBase != null) {
                        return mappedBase + "." + getterMi.getName().getIdentifier() + "()";
                    }
                }
            }
        }
        if (source != null) {
            if (constantResolver != null && context != null) {
                String resolved = constantResolver.resolve(source, context);
                if (resolved != null) {
                    return resolved;
                }
            }
            return source.toString();
        }
        return extractPayload(transformReceiver, null, constantResolver, context);
    }

    private static Expression peelJustArgument(Expression expression) {
        Expression peeled = peelFactoryPayloadExpression(expression);
        if (peeled != null) {
            return peeled;
        }
        if (expression instanceof MethodInvocation mi && mi.getExpression() != null) {
            return peelJustArgument(mi.getExpression());
        }
        return null;
    }
}
