package click.kamil.callgraphexporter;

import org.eclipse.jdt.core.JavaCore;
import org.eclipse.jdt.core.dom.*;

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.util.*;

public class JdtResolvedCallGraph {

    // Use signature as the unique key
    static Map<MethodSignature, Set<MethodSignature>> callGraph = new HashMap<>();

    // Optionally store metadata
    static Map<MethodSignature, MethodMetadata> metadataMap = new HashMap<>();

    public static void main(String[] args) throws IOException {
        File sourceDir = new File("src/main/java/click/kamil/callgraphexporter/simple");

        List<File> javaFiles = new ArrayList<>();
        collectJavaFiles(sourceDir, javaFiles);

        for (File file : javaFiles) {
            parseFileWithBinding(file, sourceDir);
        }

        generateDotFile("callgraph.dot");
        System.out.println("Call graph saved to callgraph.dot");
    }

    private static void collectJavaFiles(File dir, List<File> javaFiles) {
        for (File file : Objects.requireNonNull(dir.listFiles())) {
            if (file.isDirectory()) {
                collectJavaFiles(file, javaFiles);
            } else if (file.getName().endsWith(".java")) {
                javaFiles.add(file);
            }
        }
    }

    private static void parseFileWithBinding(File file, File sourceRoot) throws IOException {
        String source = Files.readString(file.toPath());

        ASTParser parser = ASTParser.newParser(AST.JLS21);
        parser.setSource(source.toCharArray());
        parser.setKind(ASTParser.K_COMPILATION_UNIT);
        parser.setResolveBindings(true);
        parser.setBindingsRecovery(true);

        parser.setEnvironment(
                null,
                new String[]{sourceRoot.getAbsolutePath()},
                null,
                true
        );

        parser.setUnitName(file.getName());

        Map<String, String> options = JavaCore.getOptions();
        options.put(JavaCore.COMPILER_SOURCE, JavaCore.VERSION_17);
        parser.setCompilerOptions(options);

        CompilationUnit cu = (CompilationUnit) parser.createAST(null);

        cu.accept(new ASTVisitor() {
            @Override
            public boolean visit(MethodDeclaration method) {
                IMethodBinding binding = method.resolveBinding();
                if (binding == null) return super.visit(method);

                MethodSignature callerSig = toSignature(binding);
                MethodMetadata callerMeta = toMetadata(binding, method);
                metadataMap.put(callerSig, callerMeta);
                callGraph.putIfAbsent(callerSig, new HashSet<>());

                method.accept(new ASTVisitor() {
                    @Override
                    public boolean visit(MethodInvocation node) {
                        IMethodBinding calleeBinding = node.resolveMethodBinding();
                        if (calleeBinding != null) {
                            MethodSignature calleeSig = toSignature(calleeBinding);
                            callGraph.get(callerSig).add(calleeSig);
                            // You can collect metadata for callees too if needed
                        }
                        return super.visit(node);
                    }
                });

                return super.visit(method);
            }
        });
    }

    private static MethodSignature toSignature(IMethodBinding binding) {
        String packageName = binding.getDeclaringClass().getPackage().getName();
        String className = binding.getDeclaringClass().getName();
        String methodName = binding.getName();

        List<String> paramTypes = new ArrayList<>();
        for (ITypeBinding param : binding.getParameterTypes()) {
            paramTypes.add(param.getName());
        }

        return new MethodSignature(packageName, className, methodName, paramTypes);
    }

    private static MethodMetadata toMetadata(IMethodBinding binding, MethodDeclaration method) {
        String returnType = binding.getReturnType().getName();

        Set<String> exceptions = new HashSet<>();
        for (ITypeBinding ex : binding.getExceptionTypes()) {
            exceptions.add(ex.getName());
        }

        Set<Modifiers> modifiers = extractModifiers(binding, method);
        List<String> annotations = new ArrayList<>();

        return new MethodMetadata(returnType, exceptions, modifiers, annotations, null);
    }

    private static void generateDotFile(String filename) throws IOException {
        StringBuilder sb = new StringBuilder("digraph G {\n");
        for (MethodSignature caller : callGraph.keySet()) {
            for (MethodSignature callee : callGraph.get(caller)) {
                sb.append("  \"").append(caller).append("\" -> \"").append(callee).append("\";\n");
            }
        }
        sb.append("}");
        Files.writeString(new File(filename).toPath(), sb.toString());
    }
    private static Set<Modifiers> extractModifiers(IMethodBinding binding, MethodDeclaration method) {
        Set<Modifiers> result = new HashSet<>();
        int mods = binding.getModifiers();

        // 1. From bitmask
        if (Modifier.isPublic(mods))        result.add(Modifiers.PUBLIC);
        if (Modifier.isPrivate(mods))       result.add(Modifiers.PRIVATE);
        if (Modifier.isProtected(mods))     result.add(Modifiers.PROTECTED);
        if (Modifier.isStatic(mods))        result.add(Modifiers.STATIC);
        if (Modifier.isAbstract(mods))      result.add(Modifiers.ABSTRACT);
        if (Modifier.isFinal(mods))         result.add(Modifiers.FINAL);
        if (Modifier.isNative(mods))        result.add(Modifiers.NATIVE);
        if (Modifier.isSynchronized(mods))  result.add(Modifiers.SYNCHRONIZED);
        if (Modifier.isStrictfp(mods))      result.add(Modifiers.STRICTFP);

        // 2. From modifiers list for special cases
        for (Object mod : method.modifiers()) {
            if (mod instanceof Modifier keyword) {
                if (keyword.isDefault()) result.add(Modifiers.DEFAULT);
                // NOTE: isSealed and isNonSealed are class-level only
            }
            // If needed, check for annotation modifiers here
        }

        return result;
    }
}