bridge117

package com.citi.get.cet.comet.tools;

import org.springframework.beans.PropertyValue;
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.*;

import static org.apache.commons.lang.StringUtils.trim;

public class SpringVisualizer implements ApplicationContextAware, SmartLifecycle {
    private ConfigurableApplicationContext startContext;
    private boolean isRunning = false;

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

    @Override
    public void start() {
        StringBuilder dot = new StringBuilder("digraph G {\n");
        // Use single brackets [label=<...>] for HTML rendering
        dot.append("    graph [fontname=\"Arial\", fontsize=10];\n");
        dot.append("    node [shape=plaintext, fontname=\"Arial\", fontsize=10];\n\n");

        Set<String> processed = new HashSet<>();
        ApplicationContext current = this.startContext;

        while (current != null) {
            ConfigurableListableBeanFactory factory = ((ConfigurableApplicationContext) current).getBeanFactory();
            for (String name : factory.getBeanDefinitionNames()) {
                if (processed.add(name)) {
                    BeanDefinition def = factory.getBeanDefinition(name);
                    // Logic to color-code based on your project segments
                    String color = name.toLowerCase().contains("comet") ? "#FFF904" :
                                  (name.toLowerCase().contains("citi") ? "#CBE6C7" : "#E1F5FE");

                    dot.append(String.format("    \"%s\" [label=<", name));
                    dot.append("<TABLE BORDER=\"0\" CELLBORDER=\"1\" CELLSPACING=\"0\" CELLPADDING=\"4\" BGCOLOR=\"").append(color).append("\">");

                    // Header row with fixed width and proper wrapping
                    dot.append("<TR><TD COLSPAN=\"2\" BGCOLOR=\"#DDDDDD\" CELLPADDING=\"6\"><B>").append(trim(xmlEscape(name))).append("</B></TD></TR>");

                    // 1. Process Properties (Setter injection)
                    for (PropertyValue pv : def.getPropertyValues().getPropertyValues()) {
                        dot.append("<TR>");
                        dot.append("<TD ALIGN=\"LEFT\" CELLPADDING=\"3\"><FONT FACE=\"Courier\">").append(xmlEscape(pv.getName())).append("</FONT></TD>");
                        dot.append("<TD ALIGN=\"LEFT\" CELLPADDING=\"3\">").append(trim(extractValue(pv.getValue()))).append("</TD>");
                        dot.append("</TR>");
                    }

                    // 2. Process Constructor Arguments
                    ConstructorArgumentValues cav = def.getConstructorArgumentValues();
                    // Handle Indexed (0, 1, 2...)
                    for (Map.Entry<Integer, ConstructorArgumentValues.ValueHolder> entry : cav.getIndexedArgumentValues().entrySet()) {
                        dot.append("<TR>");
                        dot.append("<TD ALIGN=\"LEFT\" CELLPADDING=\"3\"><I>arg[").append(entry.getKey()).append("]</I></TD>");
                        dot.append("<TD ALIGN=\"LEFT\" CELLPADDING=\"3\">").append(trim(extractValue(entry.getValue().getValue()))).append("</TD>");
                        dot.append("</TR>");
                    }
                    // Handle Generic/Typed
                    for (ConstructorArgumentValues.ValueHolder vh : cav.getGenericArgumentValues()) {
                        String type = vh.getType() != null ? vh.getType().substring(vh.getType().lastIndexOf('.') + 1) : "gen";
                        dot.append("<TR>");
                        dot.append("<TD ALIGN=\"LEFT\" CELLPADDING=\"3\"><I>arg:").append(type).append("</I></TD>");
                        dot.append("<TD ALIGN=\"LEFT\" CELLPADDING=\"3\">").append(trim(extractValue(vh.getValue()))).append("</TD>");
                        dot.append("</TR>");
                    }
                    dot.append("</TABLE>>];\n");

                    // 3. Draw Dependency Arrows
                    for (String dep : factory.getDependenciesForBean(name)) {
                        dot.append(String.format("    \"%s\" -> \"%s\";\n", name, dep));
                    }
                }
            }
            current = current.getParent();
        }
        dot.append("}\n");
        writeToFile(dot.toString());
        this.isRunning = true;
    }

    private String extractValue(Object value) {
        if (value == null) return "null";

        // 1. Unmask Spring's TypedStringValue immediately
        if (value instanceof TypedStringValue) {
            String raw = ((TypedStringValue) value).getValue();
            return trim(xmlEscape(raw != null ? raw.trim() : "null"));
        }

        // 2. Handle Bean References (@name)
        if (value instanceof BeanReference) {
            return "@" + trim(((BeanReference) value).getBeanName());
        }

        // 3. Handle Inner Definitions (Short Class Name only)
        if (value instanceof BeanDefinition) {
            String className = ((BeanDefinition) value).getBeanClassName();
            if (className != null) {
                return "<i>" + className.substring(className.lastIndexOf('.') + 1).trim() + "</i>";
            }
            return "<i>InnerBean</i>";
        }

        // 4. Handle Collections (Lists/Sets) - Recursive clean
        if (value instanceof Iterable) {
            List<String> cleaned = new ArrayList<>();
            int count = 0;
            for (Object item : (Iterable<?>) value) {
                if (count++ >= 15) {
                    cleaned.add("...");
                    break;
                }
                cleaned.add(trim(extractValue(item)));
            }
            return String.join("<BR ALIGN=\"LEFT\"/>", cleaned);
        }

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

        // Default fallback with aggressive trimming
        String result = value.toString().trim();
        if (result.length() > 60) {
            result = result.substring(0, 57).trim() + "...";
        }
        return trim(xmlEscape(result));
    }

    private void writeToFile(String content) {
        try {
            // Dynamically find the project root (climbing up from current dir)
            File projectRoot = findProjectRoot();
            File buildDir = new File(projectRoot, "build");
            if (!buildDir.exists()) {
                buildDir.mkdirs();
            }
            File file = new File(buildDir, "spring-beans.dot");
            try (FileWriter writer = new FileWriter(file)) {
                writer.write(content);
            }

            // Generate a proper URI for a clickable console link
            String fileUrl = file.toURI().toString();
            String sep = "==================================================";
            System.out.println("\n" + sep);
            System.out.println("GRAPH GENERATION SUCCESSFUL");
            System.out.println("Root Detected: " + projectRoot.getAbsolutePath());
            System.out.println("Click to open file: " + fileUrl);
            System.out.println(sep + "\n");
        } catch (IOException e) {
            System.err.println("CRITICAL: Failed to write DOT file: " + e.getMessage());
        }
    }

    /**
     * Climbs the directory tree to find the Gradle/Project root.
     */
    private File findProjectRoot() {
        String userDir = System.getProperty("user.dir");
        File current = new File(userDir);
        while (current != null) {
            // Look for Gradle markers to identify the project root
            if (new File(current, "build.gradle").exists() ||
                new File(current, "settings.gradle").exists() ||
                new File(current, "build.gradle.kts").exists()) {
                return current;
            }
            current = current.getParentFile();
        }
        // Fallback to current working directory if no root marker is found
        return new File(userDir);
    }

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

    @Override
    public int getPhase() { return Integer.MAX_VALUE; }

    @Override
    public boolean isAutoStartup() { return true; }

    @Override
    public void stop() { this.isRunning = false; }

    @Override
    public boolean isRunning() { return this.isRunning; }

    @Override
    public void stop(Runnable callback) { stop(); callback.run(); }
}