package click.kamil.springstatemachineexporter.analysis.flow.alloc;

import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
import org.eclipse.jdt.core.dom.ASTNode;
import org.eclipse.jdt.core.dom.ASTVisitor;
import org.eclipse.jdt.core.dom.CastExpression;
import org.eclipse.jdt.core.dom.Expression;
import org.eclipse.jdt.core.dom.IfStatement;
import org.eclipse.jdt.core.dom.InfixExpression;
import org.eclipse.jdt.core.dom.InstanceofExpression;
import org.eclipse.jdt.core.dom.MethodDeclaration;
import org.eclipse.jdt.core.dom.Modifier;
import org.eclipse.jdt.core.dom.ParenthesizedExpression;
import org.eclipse.jdt.core.dom.PatternInstanceofExpression;
import org.eclipse.jdt.core.dom.SimpleName;
import org.eclipse.jdt.core.dom.SingleVariableDeclaration;
import org.eclipse.jdt.core.dom.Statement;
import org.eclipse.jdt.core.dom.TypeDeclaration;
import org.eclipse.jdt.core.dom.TypePattern;

import java.util.ArrayList;
import java.util.List;

/**
 * Dominating {@code instanceof} / pattern / cast narrowing → unique concrete {@link AllocationSite}.
 * Shared by {@link AllocationFactPropagator}; fail-closed for else-branch, abstract targets,
 * and sends outside the narrow.
 */
public final class InstanceofNarrower {

    private InstanceofNarrower() {
    }

    /**
     * Facts when {@code useSite} is dominated by {@code recv instanceof T} (then-branch)
     * or {@code recv} is pattern variable {@code T p}.
     */
    public static ReceiverFacts factsAt(
            Expression receiver,
            ASTNode useSite,
            CodebaseContext context) {
        if (context == null || receiver == null || useSite == null) {
            return ReceiverFacts.empty();
        }
        Expression peeledRecv = peel(receiver);
        if (!(peeledRecv instanceof SimpleName recvSn)) {
            return ReceiverFacts.empty();
        }
        String recvName = recvSn.getIdentifier();
        MethodDeclaration enclosing = findEnclosingMethod(useSite);
        String patternType = resolvePatternVariableType(recvName, enclosing);
        if (patternType != null) {
            return factsFromConcreteTypeName(patternType, context);
        }
        ASTNode current = useSite;
        while (current != null && !(current instanceof MethodDeclaration)) {
            ASTNode parent = current.getParent();
            if (parent instanceof IfStatement ifStmt) {
                if (isNodeUnder(ifStmt.getThenStatement(), current)) {
                    ReceiverFacts fromCond = extractInstanceofFacts(
                            ifStmt.getExpression(), recvName, context);
                    if (!fromCond.isEmpty()) {
                        return fromCond;
                    }
                }
            }
            current = parent;
        }
        return ReceiverFacts.empty();
    }

    /**
     * Facts from a {@code (ConcreteType) expr} cast wrapping the receiver expression itself.
     */
    public static ReceiverFacts factsFromCast(Expression receiver, CodebaseContext context) {
        if (receiver == null || context == null) {
            return ReceiverFacts.empty();
        }
        Expression current = receiver;
        while (current instanceof ParenthesizedExpression pe) {
            current = pe.getExpression();
        }
        if (current instanceof CastExpression ce && ce.getType() != null) {
            return factsFromConcreteTypeName(ce.getType().toString(), context);
        }
        return ReceiverFacts.empty();
    }

    /**
     * Merge base AFP facts with dominating instanceof and/or cast.
     * Precedence: instanceof, then cast, then base. Concrete CIC conflict → empty.
     */
    public static ReceiverFacts refine(
            ReceiverFacts base,
            Expression receiver,
            ASTNode useSite,
            CodebaseContext context) {
        ReceiverFacts narrowed = factsAt(receiver, useSite, context);
        if (!narrowed.isEmpty()) {
            return mergeRefine(base, narrowed, context);
        }
        ReceiverFacts fromCast = factsFromCast(receiver, context);
        if (!fromCast.isEmpty()) {
            return mergeRefine(base, fromCast, context);
        }
        return base == null ? ReceiverFacts.empty() : base;
    }

    private static ReceiverFacts mergeRefine(
            ReceiverFacts base,
            ReceiverFacts refined,
            CodebaseContext context) {
        if (base == null || base.isEmpty()) {
            return refined;
        }
        AllocationSite baseUnique = base.uniqueOrNull();
        AllocationSite refineUnique = refined.uniqueOrNull();
        if (baseUnique == null || refineUnique == null) {
            return ReceiverFacts.empty();
        }
        if (isWeakAllocSite(baseUnique, context)) {
            return refined;
        }
        if (sameType(baseUnique.typeFqn(), refineUnique.typeFqn())) {
            return ReceiverFacts.of(baseUnique.preferRich(refineUnique));
        }
        return ReceiverFacts.empty();
    }

    public static String resolvePatternVariableType(String varName, MethodDeclaration enclosing) {
        if (varName == null || enclosing == null || enclosing.getBody() == null) {
            return null;
        }
        String[] found = {null};
        enclosing.getBody().accept(new ASTVisitor() {
            @Override
            public boolean visit(PatternInstanceofExpression node) {
                SingleVariableDeclaration svd = patternVariableOf(node);
                if (svd != null
                        && varName.equals(svd.getName().getIdentifier())
                        && svd.getType() != null) {
                    found[0] = svd.getType().toString();
                    return false;
                }
                return super.visit(node);
            }
        });
        if (found[0] == null) {
            return null;
        }
        String type = found[0];
        int generic = type.indexOf('<');
        return generic >= 0 ? type.substring(0, generic) : type;
    }

    private static boolean isWeakAllocSite(AllocationSite site, CodebaseContext context) {
        if (site == null || site.typeFqn() == null || site.typeFqn().isBlank()) {
            return true;
        }
        String fqn = site.typeFqn();
        if ("java.lang.Object".equals(fqn) || "Object".equals(fqn)) {
            return true;
        }
        TypeDeclaration td = resolveTypeDeclaration(fqn, context);
        if (td == null) {
            return true;
        }
        return td.isInterface() || Modifier.isAbstract(td.getModifiers());
    }

    private static boolean sameType(String a, String b) {
        if (a == null || b == null) {
            return false;
        }
        if (a.equals(b)) {
            return true;
        }
        String simpleA = a.contains(".") ? a.substring(a.lastIndexOf('.') + 1) : a;
        String simpleB = b.contains(".") ? b.substring(b.lastIndexOf('.') + 1) : b;
        return simpleA.equals(simpleB);
    }

    private static ReceiverFacts extractInstanceofFacts(
            Expression condition,
            String recvName,
            CodebaseContext context) {
        for (Expression part : flattenAndConditions(condition)) {
            ReceiverFacts fromPart = extractSingleInstanceof(part, recvName, context);
            if (!fromPart.isEmpty()) {
                return fromPart;
            }
        }
        return ReceiverFacts.empty();
    }

    private static List<Expression> flattenAndConditions(Expression condition) {
        List<Expression> parts = new ArrayList<>();
        collectAndParts(condition, parts);
        if (parts.isEmpty() && condition != null) {
            parts.add(condition);
        }
        return parts;
    }

    private static void collectAndParts(Expression expression, List<Expression> parts) {
        Expression cond = peel(expression);
        if (cond instanceof InfixExpression infix
                && infix.getOperator() == InfixExpression.Operator.CONDITIONAL_AND) {
            collectAndParts(infix.getLeftOperand(), parts);
            collectAndParts(infix.getRightOperand(), parts);
            return;
        }
        if (cond != null) {
            parts.add(cond);
        }
    }

    private static ReceiverFacts extractSingleInstanceof(
            Expression condition,
            String recvName,
            CodebaseContext context) {
        Expression cond = peel(condition);
        if (cond instanceof InstanceofExpression io) {
            Expression left = peel(io.getLeftOperand());
            if (left instanceof SimpleName sn && recvName.equals(sn.getIdentifier())
                    && io.getRightOperand() != null) {
                return factsFromConcreteTypeName(io.getRightOperand().toString(), context);
            }
        }
        if (cond instanceof PatternInstanceofExpression pio) {
            Expression left = peel(pio.getLeftOperand());
            if (left instanceof SimpleName sn && recvName.equals(sn.getIdentifier())) {
                SingleVariableDeclaration svd = patternVariableOf(pio);
                if (svd != null && svd.getType() != null) {
                    return factsFromConcreteTypeName(svd.getType().toString(), context);
                }
            }
        }
        return ReceiverFacts.empty();
    }

    private static ReceiverFacts factsFromConcreteTypeName(String typeFqn, CodebaseContext context) {
        if (typeFqn == null || typeFqn.isBlank()) {
            return ReceiverFacts.empty();
        }
        int generic = typeFqn.indexOf('<');
        if (generic >= 0) {
            typeFqn = typeFqn.substring(0, generic);
        }
        TypeDeclaration td = resolveTypeDeclaration(typeFqn, context);
        if (td == null || td.isInterface() || Modifier.isAbstract(td.getModifiers())) {
            return ReceiverFacts.empty();
        }
        String fqn = context.getFqn(td);
        return ReceiverFacts.of(new AllocationSite(fqn != null ? fqn : typeFqn));
    }

    private static TypeDeclaration resolveTypeDeclaration(String typeFqn, CodebaseContext context) {
        TypeDeclaration td = context.getTypeDeclaration(typeFqn);
        if (td == null) {
            String simple = typeFqn.contains(".") ? typeFqn.substring(typeFqn.lastIndexOf('.') + 1) : typeFqn;
            td = context.getTypeDeclaration(simple);
        }
        return td;
    }

    private static SingleVariableDeclaration patternVariableOf(PatternInstanceofExpression node) {
        if (node == null) {
            return null;
        }
        try {
            if (node.getPattern() instanceof TypePattern tp) {
                SingleVariableDeclaration pv = tp.getPatternVariable();
                if (pv != null) {
                    return pv;
                }
            }
        } catch (RuntimeException ignored) {
            // API level may expose rightOperand instead of pattern
        }
        try {
            return node.getRightOperand();
        } catch (RuntimeException ignored) {
            return null;
        }
    }

    private static boolean isNodeUnder(Statement root, ASTNode node) {
        if (root == null || node == null) {
            return false;
        }
        ASTNode current = node;
        while (current != null) {
            if (current == root) {
                return true;
            }
            if (current instanceof MethodDeclaration) {
                return false;
            }
            current = current.getParent();
        }
        return false;
    }

    private static MethodDeclaration findEnclosingMethod(ASTNode node) {
        ASTNode current = node;
        while (current != null) {
            if (current instanceof MethodDeclaration md) {
                return md;
            }
            current = current.getParent();
        }
        return null;
    }

    private static Expression peel(Expression expression) {
        Expression current = expression;
        while (current instanceof ParenthesizedExpression pe) {
            current = pe.getExpression();
        }
        if (current instanceof CastExpression ce) {
            return peel(ce.getExpression());
        }
        return current;
    }
}
