bridge104

package com.yourpackage;

import org.springframework.beans.PropertyValue;
import org.springframework.beans.factory.BeanDefinitionStoreException;
import org.springframework.beans.factory.config.*;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.ConfigurableApplicationContext;
import org.springframework.context.SmartLifecycle;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.*;

public class SpringVisualizer implements ApplicationContextAware, SmartLifecycle {

    private ConfigurableApplicationContext context;
    private boolean isRunning = false;
    private final Set<String> currentBeanManualDeps = new HashSet<>();

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) {
        this.context = (ConfigurableApplicationContext) applicationContext;
    }

    @Override
    public void start() {
        StringBuilder dot = new StringBuilder("digraph G {\n");
        dot.append("  rankdir=LR; node [shape=plain, fontname=\"Arial\", fontsize=10];\n");
        dot.append("  edge [fontname=\"Arial\", fontsize=8];\n\n");

        Set<String> processed = new HashSet<>();

        // Find the absolute root of the context tree to ensure we see everything
        ApplicationContext root = context;
        while (root.getParent() != null) {
            root = root.getParent();
        }

        // Use a Queue to traverse all child contexts if they exist
        Queue<ApplicationContext> queue = new LinkedList<>();
        queue.add(root);

        while (!queue.isEmpty()) {
            ApplicationContext ctx = queue.poll();
            if (!(ctx instanceof ConfigurableApplicationContext)) continue;

            ConfigurableListableBeanFactory factory = ((ConfigurableApplicationContext) ctx).getBeanFactory();

            for (String name : factory.getBeanDefinitionNames()) {
                if (!processed.add(name)) continue;

                try {
                    BeanDefinition def = factory.getBeanDefinition(name);
                    currentBeanManualDeps.clear();

                    String color = name.toLowerCase().contains("comet") ? "#FFF9C4" : 
                                   (name.toLowerCase().contains("cpls") ? "#C8E6C9" : "#E1F5FE");

                    dot.append(String.format("  \"%s\" [label=<", name));
                    dot.append("<TABLE BORDER=\"0\" CELLBORDER=\"1\" CELLSPACING=\"0\" BGCOLOR=\"").append(color).append("\">");
                    dot.append("<TR><TD COLSPAN=\"2\"><B>").append(xmlEscape(name.trim())).append("</B></TD></TR>");

                    for (PropertyValue pv : def.getPropertyValues().getPropertyValues()) {
                        dot.append("<TR><TD ALIGN=\"LEFT\">").append(xmlEscape(pv.getName().trim())).append("</TD>");
                        dot.append("<TD ALIGN=\"LEFT\">").append(extractValue(pv.getValue(), 0)).append("</TD></TR>");
                    }

                    ConstructorArgumentValues cav = def.getConstructorArgumentValues();
                    for (Map.Entry<Integer, ConstructorArgumentValues.ValueHolder> entry : cav.getIndexedArgumentValues().entrySet()) {
                        dot.append("<TR><TD ALIGN=\"LEFT\"><I>arg[").append(entry.getKey()).append("]</I></TD>");
                        dot.append("<TD ALIGN=\"LEFT\">").append(extractValue(entry.getValue().getValue(), 0)).append("</TD></TR>");
                    }
                    for (ConstructorArgumentValues.ValueHolder vh : cav.getGenericArgumentValues()) {
                        dot.append("<TR><TD ALIGN=\"LEFT\"><I>arg:gen</I></TD>");
                        dot.append("<TD ALIGN=\"LEFT\">").append(extractValue(vh.getValue(), 0)).append("</TD></TR>");
                    }
                    dot.append("</TABLE>>];\n");

                    // Standard + Manual Dependencies
                    Set<String> allDeps = new HashSet<>(Arrays.asList(factory.getDependenciesForBean(name)));
                    allDeps.addAll(currentBeanManualDeps);
                    for (String dep : allDeps) {
                        dot.append(String.format("  \"%s\" -> \"%s\";\n", name, dep));
                    }
                } catch (BeanDefinitionStoreException ignored) {
                    // Some beans in parent/child hierarchies might be inaccessible
                }
            }

            // Add children to queue (if any)
            // Note: Standard ApplicationContext doesn't expose children easily, 
            // but starting from Root usually covers the hierarchy via getBeanDefinitionNames()
        }

        dot.append("}\n");
        writeToFile(dot.toString());
        this.isRunning = true;
    }

    private String extractValue(Object value, int depth) {
        if (value == null) return "null";
        if (depth > 8) return "<i>[Max Depth]</i>";

        if (value instanceof BeanReference) {
            String beanName = ((BeanReference) value).getBeanName().trim();
            currentBeanManualDeps.add(beanName);
            return "@" + xmlEscape(beanName);
        }

        if (value instanceof TypedStringValue) {
            String raw = ((TypedStringValue) value).getValue();
            return xmlEscape(raw != null ? raw.trim() : "null");
        }

        // FIX: Explicitly handle both BeanDefinition and BeanDefinitionHolder
        if (value instanceof BeanDefinition || value instanceof BeanDefinitionHolder) {
            BeanDefinition innerDef = (value instanceof BeanDefinitionHolder) ? 
                                       ((BeanDefinitionHolder) value).getBeanDefinition() : (BeanDefinition) value;

            StringBuilder innerHtml = new StringBuilder("<TABLE BORDER=\"0\" CELLBORDER=\"1\" CELLSPACING=\"0\" BGCOLOR=\"#F5F5F5\">");
            String className = innerDef.getBeanClassName();
            String shortName = className != null ? className.substring(className.lastIndexOf('.') + 1) : 
                               (innerDef.getParentName() != null ? "Parent:" + innerDef.getParentName() : "InnerBean");

            innerHtml.append("<TR><TD COLSPAN=\"2\" BGCOLOR=\"#DCDCDC\"><B>").append(xmlEscape(shortName.trim())).append("</B></TD></TR>");

            for (PropertyValue pv : innerDef.getPropertyValues().getPropertyValues()) {
                innerHtml.append("<TR><TD><FONT POINT-SIZE=\"8\">").append(xmlEscape(pv.getName().trim())).append("</FONT></TD>");
                innerHtml.append("<TD>").append(extractValue(pv.getValue(), depth + 1)).append("</TD></TR>");
            }
            innerHtml.append("</TABLE>");
            return innerHtml.toString();
        }

        if (value instanceof Iterable) {
            List<String> cleaned = new ArrayList<>();
            for (Object item : (Iterable<?>) value) {
                cleaned.add(extractValue(item, depth + 1));
            }
            return String.join("<BR ALIGN=\"LEFT\"/>", cleaned);
        }

        if (value instanceof Map) {
            StringBuilder mapStr = new StringBuilder();
            for (Map.Entry<?, ?> entry : ((Map<?, ?>) value).entrySet()) {
                mapStr.append(extractValue(entry.getKey(), depth + 1)).append("=").append(extractValue(entry.getValue(), depth + 1)).append("<BR ALIGN=\"LEFT\"/>");
            }
            return mapStr.toString();
        }

        String s = value.toString().trim();
        return xmlEscape(s.length() > 60 ? s.substring(0, 57).trim() + "..." : s);
    }

    private void writeToFile(String content) {
        try {
            File root = findProjectRoot();
            File buildDir = new File(root, "build");
            if (!buildDir.exists()) buildDir.mkdirs();
            File file = new File(buildDir, "spring-beans.dot");
            try (FileWriter writer = new FileWriter(file)) {
                writer.write(content);
            }
            System.out.println("\n============================================================");
            System.out.println("GRAPH GENERATION SUCCESSFUL");
            System.out.println("URL: " + file.toURI().toString());
            System.out.println("============================================================\n");
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private File findProjectRoot() {
        File current = new File(System.getProperty("user.dir"));
        while (current != null) {
            if (new File(current, "build.gradle").exists() || new File(current, "settings.gradle").exists()) return current;
            current = current.getParentFile();
        }
        return new File(System.getProperty("user.dir"));
    }

    private String xmlEscape(String input) {
        if (input == null) return "";
        return input.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
                    .replace("\"", "&quot;").replace("'", "&apos;")
                    .replace("&lt;BR ALIGN=&quot;LEFT&quot;/&gt;", "<BR ALIGN=\"LEFT\"/>").trim();
    }

    @Override public int getPhase() { return Integer.MAX_VALUE; }
    @Override public boolean isAutoStartup() { return true; }
    @Override public void stop() { isRunning = false; }
    @Override public boolean isRunning() { return isRunning; }
    @Override public void stop(Runnable c) { stop(); c.run(); }
}