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

import click.kamil.springstatemachineexporter.analysis.pipeline.AccessorResolver;
import click.kamil.springstatemachineexporter.analysis.pipeline.ResolutionBudget;
import click.kamil.springstatemachineexporter.analysis.resolver.ConstantResolver;
import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
import org.eclipse.jdt.core.dom.CompilationUnit;
import org.eclipse.jdt.core.dom.Expression;
import org.eclipse.jdt.core.dom.MethodDeclaration;
import org.eclipse.jdt.core.dom.ReturnStatement;
import org.eclipse.jdt.core.dom.TypeDeclaration;
import org.eclipse.jdt.core.dom.ASTVisitor;

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

/**
 * Given concrete allocation facts and an accessor name ({@code getType}/{@code getEvent}/…),
 * resolve overrides and extract compile-time return constants. Succeeds only when exactly one
 * constant is proven across all facts.
 */
public final class VirtualAccessorConstResolver {

    private final CodebaseContext context;
    private final AccessorResolver accessorResolver;
    private final ConstantResolver constantResolver;

    public VirtualAccessorConstResolver(
            CodebaseContext context,
            AccessorResolver accessorResolver,
            ConstantResolver constantResolver) {
        this.context = context;
        this.accessorResolver = accessorResolver;
        this.constantResolver = constantResolver;
    }

    /**
     * @return singleton list with the proven constant, or empty when unknown / ambiguous / non-constant
     */
    public List<String> resolveUniqueConstant(ReceiverFacts facts, String accessorName) {
        if (facts == null || facts.isEmpty() || accessorName == null || accessorName.isBlank()) {
            return List.of();
        }
        LinkedHashSet<String> constants = new LinkedHashSet<>();
        for (AllocationSite site : facts.sites()) {
            List<String> fromSite = constantsForSite(site, accessorName);
            if (fromSite.isEmpty()) {
                // Any site without a compile-time constant → fail closed.
                return List.of();
            }
            constants.addAll(fromSite);
            if (constants.size() > 1) {
                return List.of();
            }
        }
        if (constants.size() != 1) {
            return List.of();
        }
        return List.copyOf(constants);
    }

    private List<String> constantsForSite(AllocationSite site, String accessorName) {
        if (site == null || site.typeFqn() == null) {
            return List.of();
        }
        String typeFqn = resolveTypeFqn(site.typeFqn());
        if (typeFqn == null) {
            return List.of();
        }
        TypeDeclaration td = context.getTypeDeclaration(typeFqn);
        CompilationUnit cu = td != null && td.getRoot() instanceof CompilationUnit c ? c : null;

        List<String> viaAccessor = List.of();
        if (accessorResolver != null) {
            viaAccessor = accessorResolver.resolve(
                    typeFqn,
                    accessorName,
                    new AccessorResolver.ReceiverContext(td, null, cu, false, null),
                    ResolutionBudget.defaults());
        }
        LinkedHashSet<String> cleaned = new LinkedHashSet<>();
        for (String value : viaAccessor) {
            String normalized = normalizeConstant(value);
            if (normalized != null) {
                cleaned.add(normalized);
            }
        }
        if (cleaned.size() == 1) {
            return List.copyOf(cleaned);
        }
        if (cleaned.size() > 1) {
            return List.of();
        }

        // Fallback: evaluate return expressions with ConstantResolver (handles Enum.CONST.name()).
        if (td == null || constantResolver == null) {
            return List.of();
        }
        MethodDeclaration method = context.findMethodDeclaration(td, accessorName, true);
        if (method == null || method.getBody() == null) {
            return List.of();
        }
        List<String> fromReturns = new ArrayList<>();
        method.getBody().accept(new ASTVisitor() {
            @Override
            public boolean visit(ReturnStatement node) {
                Expression expr = node.getExpression();
                if (expr == null) {
                    return super.visit(node);
                }
                String resolved = constantResolver.resolve(expr, context);
                String normalized = normalizeConstant(resolved);
                if (normalized != null) {
                    fromReturns.add(normalized);
                }
                return super.visit(node);
            }
        });
        LinkedHashSet<String> unique = new LinkedHashSet<>(fromReturns);
        if (unique.size() == 1) {
            return List.copyOf(unique);
        }
        return List.of();
    }

    private String resolveTypeFqn(String typeName) {
        if (typeName == null || typeName.isBlank()) {
            return null;
        }
        if (typeName.contains(".")) {
            if (context.getTypeDeclaration(typeName) != null) {
                return typeName;
            }
            // May be simple-with-package missing; still try as-is for resolve.
            return typeName;
        }
        TypeDeclaration td = context.getTypeDeclaration(typeName);
        if (td != null) {
            String fqn = context.getFqn(td);
            return fqn != null ? fqn : typeName;
        }
        return typeName;
    }

    /**
     * Drop symbolic / enum-set / blank; keep FQN or bare constant identifiers.
     */
    public static String normalizeConstant(String value) {
        if (value == null || value.isBlank()) {
            return null;
        }
        if (value.startsWith("<SYMBOLIC:") || value.startsWith("ENUM_SET:") || value.contains(".*>")) {
            return null;
        }
        String trimmed = value.trim();
        // Enum.CONST.name() peel may already be bare; QualName may be FQN.
        return trimmed;
    }
}
