package click.kamil.springstatemachineexporter.ast.app;

import click.kamil.springstatemachineexporter.analysis.resolver.ConstantResolver;
import click.kamil.springstatemachineexporter.model.Action;
import click.kamil.springstatemachineexporter.model.Guard;
import click.kamil.springstatemachineexporter.model.State;
import click.kamil.springstatemachineexporter.model.Event;
import click.kamil.springstatemachineexporter.model.Transition;
import click.kamil.springstatemachineexporter.model.TransitionType;
import click.kamil.springstatemachineexporter.ast.common.AstUtils;
import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
import org.eclipse.jdt.core.dom.*;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;
import lombok.extern.slf4j.Slf4j;

@Slf4j
public class AstTransitionParser {
    private static final ConstantResolver constantResolver = new ConstantResolver();

    private static final Set<String> SUPPORTED_TRANSITION_METHOD_NAMES = Arrays.stream(TransitionType.values())
            .map(TransitionType::getMethodName)
            .collect(Collectors.toSet());

    public static List<Transition> parseTransitions(MethodDeclaration method, CodebaseContext context) {
        if (method == null)
            return Collections.emptyList();
        List<Transition> transitions = new ArrayList<>();

        method.accept(new ASTVisitor() {
            @Override
            public boolean visit(MethodInvocation node) {
                if (isEntryPoint(node)) {
                    transitions.addAll(parseTransitionsFromExpression(node, new HashMap<>(), false, context));
                }
                return super.visit(node);
            }

            private boolean isEntryPoint(MethodInvocation node) {
                // Check if this call is part of a larger chain already handled
                if (node.getParent() instanceof MethodInvocation) {
                    return false;
                }
                List<MethodInvocation> chain = collectMethodInvocationChain(node);
                return chain.stream().anyMatch(mi -> SUPPORTED_TRANSITION_METHOD_NAMES.contains(mi.getName().getIdentifier()));
            }
        });

        return transitions;
    }

    private static List<MethodInvocation> collectMethodInvocationChain(MethodInvocation root) {
        List<MethodInvocation> chain = new ArrayList<>();

        MethodInvocation current = root;
        while (current != null) {
            chain.add(0, current); // add to the front to preserve order

            Expression expr = current.getExpression();
            if (expr instanceof MethodInvocation parent) {
                current = parent;
            } else {
                break;
            }
        }

        return chain;
    }

    public static List<Transition> parseTransitionsFromExpression(MethodInvocation rootCall, Map<String, Expression> argsMap, CodebaseContext context) {
        return parseTransitionsFromExpression(rootCall, argsMap, false, context);
    }

    public static List<Transition> parseTransitionsFromExpression(MethodInvocation rootCall, Map<String, Expression> argsMap, boolean isContinuation, CodebaseContext context) {
        List<Transition> transitions = new ArrayList<>();
        CompilationUnit cu = (CompilationUnit) rootCall.getRoot();

        // Step 1: Unroll method chain
        List<MethodInvocation> calls = new ArrayList<>();
        Expression current = rootCall;
        while (current instanceof MethodInvocation mi) {
            calls.addFirst(mi);
            current = mi.getExpression();
        }

        // Step 2: Split by .and()
        List<List<MethodInvocation>> segments = new ArrayList<>();
        List<MethodInvocation> currentSegment = new ArrayList<>();
        for (MethodInvocation call : calls) {
            if ("and".equals(call.getName().getIdentifier())) {
                if (!currentSegment.isEmpty()) {
                    segments.add(currentSegment);
                    currentSegment = new ArrayList<>();
                }
            } else {
                currentSegment.add(call);
            }
        }
        if (!currentSegment.isEmpty())
            segments.add(currentSegment);

        // Step 3: Parse each segment
        for (int i = 0; i < segments.size(); i++) {
            List<MethodInvocation> segment = segments.get(i);

            Optional<TransitionType> segmentType = segment.stream()
                    .map(mi -> TransitionType.fromMethodName(mi.getName().getIdentifier()))
                    .flatMap(Optional::stream)
                    .findFirst();

            // If it's a continuation and we are at the first segment, we should be careful
            boolean isFirstSegmentOfContinuation = isContinuation && i == 0;

            if (segmentType.isPresent() && (segmentType.get() == TransitionType.CHOICE || segmentType.get() == TransitionType.JUNCTION)) {
                List<Transition> choiceTransitions = parseChoiceOrJunctionTransitions(segment, segmentType.get(), argsMap, context, cu);
                transitions.addAll(choiceTransitions);
            } else {
                Transition t = parseStandardTransition(segment, argsMap, context, cu);
                // If it's a continuation's first segment and it has no source, it's just refining the receiver
                if (isFirstSegmentOfContinuation && t.getSourceStates().isEmpty()) {
                    // Don't add it as a new transition if it doesn't have source/target
                    if (!t.getTargetStates().isEmpty()) {
                        transitions.add(t);
                    }
                } else if (!t.getSourceStates().isEmpty() || !t.getTargetStates().isEmpty() || t.getType() != null) {
                    transitions.add(t);
                }
            }
        }

        return transitions;
    }

    private static List<Transition> parseChoiceOrJunctionTransitions(List<MethodInvocation> segment, TransitionType type, Map<String, Expression> argsMap, CodebaseContext context, CompilationUnit cu) {
        List<Transition> transitions = new ArrayList<>();
        State sourceState = null;
        int orderCounter = 0;

        // Extract source state once
        for (MethodInvocation call : segment) {
            if ("source".equals(call.getName().getIdentifier()) && !call.arguments().isEmpty()) {
                Expression arg = resolveArg((Expression) call.arguments().get(0), argsMap);
                sourceState = context.resolveState(arg, cu);
                break; // only take the first occurrence
            }
        }

        for (MethodInvocation call : segment) {
            String methodName = call.getName().getIdentifier();
            List<?> args = call.arguments();
            if (!List.of("first", "then", "last").contains(methodName) || args.isEmpty()) {
                continue;
            }
            Transition t = new Transition();
            t.setType(type);
            if (sourceState != null) {
                t.getSourceStates().add(sourceState);
            }

            Expression targetArg = resolveArg((Expression) args.get(0), argsMap);
            t.getTargetStates().add(context.resolveState(targetArg, cu));

            if (args.size() > 1) {
                Expression secondArg = resolveArg((Expression) args.get(1), argsMap);
                if ("last".equals(methodName)) {
                    // For 'last', the second argument is an action, not a guard
                    parseAction(secondArg, t, cu, context);
                } else {
                    // For 'first' and 'then', the second argument is a guard
                    parseGuard(secondArg, t, cu, context);
                }
            }
            if (args.size() > 2) {
                Expression thirdArg = resolveArg((Expression) args.get(2), argsMap);
                parseAction(thirdArg, t, cu, context);
            }
            t.setOrder(orderCounter++);
            transitions.add(t);
        }
        return transitions;
    }

    private static Transition parseStandardTransition(List<MethodInvocation> segment, Map<String, Expression> argsMap, CodebaseContext context, CompilationUnit cu) {
        Transition t = new Transition();
        for (MethodInvocation call : segment) {
            String methodName = call.getName().getIdentifier();
            List<?> args = call.arguments();

            Optional<TransitionType> type = TransitionType.fromMethodName(methodName);
            if (type.isPresent()) {
                t.setType(type.get());
                continue;
            }
            if (args.isEmpty()) {
                continue;
            }

            switch (methodName) {
                case "source" -> args.forEach(arg -> {
                    Expression resolved = resolveArg((Expression) arg, argsMap);
                    t.getSourceStates().add(context.resolveState(resolved, cu));
                });
                case "target" -> args.forEach(arg -> {
                    Expression resolved = resolveArg((Expression) arg, argsMap);
                    t.getTargetStates().add(context.resolveState(resolved, cu));
                });
                 case "event" -> {
                    Expression resolved = resolveArg((Expression) args.get(0), argsMap);
                    String eventValue = constantResolver.resolve(resolved, context);
                    if (eventValue != null) {
                        t.setEvent(Event.of(resolved.toString(), eventValue));
                    } else if (resolved instanceof QualifiedName qn) {
                        String qualifier = qn.getQualifier().toString();
                        String name = qn.getName().getIdentifier();
                        TypeDeclaration td = context.getTypeDeclaration(qualifier, cu);
                        if (td != null) {
                            t.setEvent(Event.of(resolved.toString(), context.getFqn(td) + "." + name));
                        } else {
                            t.setEvent(Event.of(resolved.toString(), qn.getFullyQualifiedName()));
                        }
                    } else if (resolved instanceof SimpleName sn) {
                        String name = sn.getIdentifier();
                        String full = name;
                        for (Object impObj : cu.imports()) {
                            ImportDeclaration imp = (ImportDeclaration) impObj;
                            String impName = imp.getName().getFullyQualifiedName();
                            if (impName.endsWith("." + name)) {
                                full = impName;
                                break;
                            }
                        }
                        t.setEvent(Event.of(resolved.toString(), full));
                    } else {
                        t.setEvent(Event.of(QuotedExpression.of(resolved).toStringWithoutQuotes()));
                    }
                    if (args.size() > 1) {
                        Expression actionArg = resolveArg((Expression) args.get(1), argsMap);
                        parseAction(actionArg, t, cu, context);
                    }
                }
                case "guard", "guardExpression" -> {
                    Expression resolved = resolveArg((Expression) args.get(0), argsMap);
                    parseGuard(resolved, t, cu, context);
                }
                case "guards" -> {
                    args.forEach(arg -> {
                        Expression resolved = resolveArg((Expression) arg, argsMap);
                        parseGuard(resolved, t, cu, context);
                    });
                }
                case "action" -> {
                    Expression resolved = resolveArg((Expression) args.get(0), argsMap);
                    parseAction(resolved, t, cu, context);
                }
                case "actions" -> {
                    args.forEach(arg -> {
                        Expression resolved = resolveArg((Expression) arg, argsMap);
                        parseAction(resolved, t, cu, context);
                    });
                }
                case "timer", "timerOnce" -> {
                    Expression resolved = resolveArg((Expression) args.get(0), argsMap);
                    t.setEvent(Event.of(methodName + "(" + resolved.toString() + ")", methodName + "(" + resolved.toString() + ")"));
                }
            }
        }
        return t;
    }

    private static String extractInternalLogic(Expression expr, CompilationUnit cu, CodebaseContext context) {
        if (expr == null) return null;

        // 1. Lambda Expressions
        if (expr instanceof LambdaExpression le) {
            return le.toString();
        }

        // 2. Anonymous Classes
        if (expr instanceof ClassInstanceCreation cic && cic.getAnonymousClassDeclaration() != null) {
            return cic.toString();
        }

        // 3. Method Invocations
        if (expr instanceof MethodInvocation mi) {
            IMethodBinding binding = mi.resolveMethodBinding();
            if (binding != null) {
                // Try to resolve deeper logic from the return type (e.g., a Bean returning an Action interface)
                ITypeBinding returnType = binding.getReturnType();
                if (returnType != null) {
                    String returnFqn = returnType.getTypeDeclaration().getQualifiedName();
                    TypeDeclaration returnTd = context.getTypeDeclaration(returnFqn);
                    if (returnTd != null) {
                        for (MethodDeclaration rmd : returnTd.getMethods()) {
                            String name = rmd.getName().getIdentifier();
                            if (name.equals("execute") || name.equals("evaluate")) {
                                return rmd.toString();
                            }
                        }
                    }
                }

                // Try same file
                ASTNode declNode = cu.findDeclaringNode(binding.getKey());
                if (declNode instanceof MethodDeclaration md) {
                    if (md.getBody() != null && md.getBody().statements().size() == 1) {
                        Object stmt = md.getBody().statements().get(0);
                        if (stmt instanceof org.eclipse.jdt.core.dom.ReturnStatement rs) {
                            Expression retExpr = rs.getExpression();
                            if (isLambdaOrAnonymous(retExpr)) {
                                return retExpr.toString();
                            }
                        }
                    }
                    return md.toString();
                }

                // Try other files via CodebaseContext
                ITypeBinding declaringClass = binding.getDeclaringClass();
                if (declaringClass != null) {
                    String fqn = declaringClass.getTypeDeclaration().getQualifiedName();
                    TypeDeclaration td = context.getTypeDeclaration(fqn);
                    if (td != null) {
                        for (MethodDeclaration md : td.getMethods()) {
                            if (md.getName().getIdentifier().equals(binding.getName()) &&
                                    md.parameters().size() == binding.getParameterTypes().length) {
                                
                                // Special case: if it's a getter/factory for Action/Guard, look into the return statement
                                if (md.getBody() != null) {
                                    if (md.getBody().statements().size() == 1) {
                                        Object stmt = md.getBody().statements().get(0);
                                        if (stmt instanceof org.eclipse.jdt.core.dom.ReturnStatement rs) {
                                            Expression retExpr = rs.getExpression();
                                            if (isLambdaOrAnonymous(retExpr)) {
                                                return retExpr.toString();
                                            }
                                        }
                                    }
                                }
                                return md.toString();
                            }
                        }
                    }
                }
            } else {
                // FALLBACK: Manual resolution by name if binding failed
                TypeDeclaration currentClass = AstUtils.findEnclosingType(mi);
                if (currentClass != null) {
                    return resolveMethodManually(mi.getName().getIdentifier(), currentClass, context);
                }
            }
        }

        // 4. Method References
        if (expr instanceof MethodReference mr) {
            IMethodBinding binding = mr.resolveMethodBinding();
            if (binding != null) {
                ITypeBinding declaringClass = binding.getDeclaringClass();
                if (declaringClass != null) {
                    String fqn = declaringClass.getTypeDeclaration().getQualifiedName();
                    TypeDeclaration td = context.getTypeDeclaration(fqn);
                    if (td != null) {
                        for (MethodDeclaration md : td.getMethods()) {
                            if (md.getName().getIdentifier().equals(binding.getName()) &&
                                    md.parameters().size() == binding.getParameterTypes().length) {
                                return md.toString();
                            }
                        }
                    }
                }
            }
            return mr.toString();
        }

        // 5. Bean Reference (SimpleName, FieldAccess, etc.)
        String resolvedBeanBody = resolveBeanClassBody(expr, context);
        if (resolvedBeanBody != null) {
            return resolvedBeanBody;
        }

        return null;
    }

    private static String resolveBeanClassBody(Expression expr, CodebaseContext context) {
        String typeFqn = null;
        ITypeBinding binding = expr.resolveTypeBinding();
        if (binding != null) {
            typeFqn = binding.getQualifiedName();
        } else if (expr instanceof SimpleName sn) {
            // Check enclosing method first for local variables
            MethodDeclaration enclosingMethod = AstUtils.findEnclosingMethod(expr);
            if (enclosingMethod != null) {
                final String[] foundType = {null};
                final Expression[] foundInitializer = {null};
                enclosingMethod.accept(new ASTVisitor() {
                    @Override
                    public boolean visit(VariableDeclarationStatement vds) {
                        for (Object fragObj : vds.fragments()) {
                            if (fragObj instanceof VariableDeclarationFragment fragment) {
                                if (fragment.getName().getIdentifier().equals(sn.getIdentifier())) {
                                    foundType[0] = vds.getType().toString();
                                    foundInitializer[0] = fragment.getInitializer();
                                }
                            }
                        }
                        return super.visit(vds);
                    }
                });
                if (foundInitializer[0] != null) {
                    ASTNode root = expr.getRoot();
                    if (root instanceof CompilationUnit cu) {
                        String initLogic = extractInternalLogic(foundInitializer[0], cu, context);
                        if (initLogic != null) return initLogic;
                    }
                }
                if (foundType[0] != null) {
                    typeFqn = foundType[0];
                }
            }

            if (typeFqn == null) {
                TypeDeclaration enclosingClass = AstUtils.findEnclosingType(expr);
                if (enclosingClass != null) {
                    for (FieldDeclaration fd : enclosingClass.getFields()) {
                        for (Object fragObj : fd.fragments()) {
                            if (fragObj instanceof VariableDeclarationFragment fragment) {
                                if (fragment.getName().getIdentifier().equals(sn.getIdentifier())) {
                                    if (fragment.getInitializer() != null) {
                                        ASTNode root = expr.getRoot();
                                        if (root instanceof CompilationUnit cu) {
                                            String initLogic = extractInternalLogic(fragment.getInitializer(), cu, context);
                                            if (initLogic != null) return initLogic;
                                        }
                                    }
                                    typeFqn = fd.getType().toString();
                                    break;
                                }
                            }
                        }
                    }
                }
            }
        }

        if (typeFqn != null) {
            if (typeFqn.contains("<")) {
                typeFqn = typeFqn.substring(0, typeFqn.indexOf('<'));
            }
            TypeDeclaration beanTd = context.getTypeDeclaration(typeFqn);
            if (beanTd == null) {
                ASTNode root = expr.getRoot();
                if (root instanceof CompilationUnit cu) {
                    for (Object impObj : cu.imports()) {
                        ImportDeclaration imp = (ImportDeclaration) impObj;
                        if (!imp.isStatic() && !imp.isOnDemand()) {
                            String impName = imp.getName().getFullyQualifiedName();
                            if (impName.endsWith("." + typeFqn)) {
                                beanTd = context.getTypeDeclaration(impName);
                                break;
                            }
                        }
                    }
                }
            }
            if (beanTd != null) {
                for (MethodDeclaration md : beanTd.getMethods()) {
                    String name = md.getName().getIdentifier();
                    if ("evaluate".equals(name) || "execute".equals(name)) {
                        return md.toString();
                    }
                }
                return beanTd.toString();
            }
        }
        return null;
    }

    private static String resolveMethodManually(String methodName, TypeDeclaration td, CodebaseContext context) {
        // Search in this class
        for (MethodDeclaration md : td.getMethods()) {
            if (md.getName().getIdentifier().equals(methodName)) {
                // Found it. Apply the return-statement-extraction logic here too
                if (md.getBody() != null && md.getBody().statements().size() == 1) {
                    Object stmt = md.getBody().statements().get(0);
                    if (stmt instanceof org.eclipse.jdt.core.dom.ReturnStatement rs) {
                        Expression retExpr = rs.getExpression();
                        if (isLambdaOrAnonymous(retExpr)) {
                            return retExpr.toString();
                        }
                    }
                }
                return md.toString();
            }
        }
        
        // Search in superclass
        if (td.getSuperclassType() != null) {
            String superName = td.getSuperclassType().toString();
            // Try to resolve simple name to FQN if possible, or just search by simple name
            TypeDeclaration superTd = context.getTypeDeclaration(superName);
            if (superTd != null) {
                return resolveMethodManually(methodName, superTd, context);
            }
        }
        
        return null;
    }

    private static void parseGuard(Object arg, Transition t, CompilationUnit cu, CodebaseContext context) {
        QuotedExpression quotedExpr = QuotedExpression.of(arg);
        if (quotedExpr != null) {
            String internalLogic = extractInternalLogic(quotedExpr.getExpression(), cu, context);
            int lineNumber = cu.getLineNumber(quotedExpr.getExpression().getStartPosition());
            TypeDeclaration td = AstUtils.findEnclosingType(quotedExpr.getExpression());
            String fqn = td != null ? context.getFqn(td) : null;
            String sourceFile = fqn != null ? context.getRelativePath(fqn) : null;
            t.getGuards().add(Guard.of(quotedExpr.toStringWithoutQuotes(), isLambdaOrAnonymous(quotedExpr.getExpression()), internalLogic, lineNumber, fqn, sourceFile));
        }
    }

    private static void parseAction(Object arg, Transition t, CompilationUnit cu, CodebaseContext context) {
        QuotedExpression quotedExpr = QuotedExpression.of(arg);
        if (quotedExpr != null) {
            String internalLogic = extractInternalLogic(quotedExpr.getExpression(), cu, context);
            int lineNumber = cu.getLineNumber(quotedExpr.getExpression().getStartPosition());
            TypeDeclaration td = AstUtils.findEnclosingType(quotedExpr.getExpression());
            String fqn = td != null ? context.getFqn(td) : null;
            String sourceFile = fqn != null ? context.getRelativePath(fqn) : null;
            t.getActions().add(Action.of(quotedExpr.toStringWithoutQuotes(), isLambdaOrAnonymous(quotedExpr.getExpression()), internalLogic, lineNumber, fqn, sourceFile));
        }
    }

    private static Expression resolveArg(Expression expr, Map<String, Expression> argsMap) {
        if (expr instanceof SimpleName sn) {
            Expression resolved = argsMap.get(sn.getIdentifier());
            if (resolved != null)
                return resolved;
        }
        return expr;
    }

    public static boolean isLambdaOrAnonymous(Expression expr) {
        return switch (expr) {
            case null -> false;
            case LambdaExpression le -> true;
            case ClassInstanceCreation cic -> cic.getAnonymousClassDeclaration() != null;
            default -> false;
        };
    }
}
