package click.kamil.springstatemachineexporter.exporter;

import click.kamil.springstatemachineexporter.model.State;
import click.kamil.springstatemachineexporter.model.Transition;
import click.kamil.springstatemachineexporter.model.TransitionType;

import java.util.HashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

public class Dot implements StateMachineExporter {


    private static final String[] CHOICE_COLORS = {
            "blue", "green", "red", "purple", "orange", "brown", "darkgreen", "darkblue"
    };

    @Override
    public String export(click.kamil.springstatemachineexporter.analysis.model.AnalysisResult result, ExportOptions options) {
        StringBuilder sb = new StringBuilder();
        sb.append("digraph statemachine {\n");
        sb.append("  rankdir=LR;\n");
        sb.append("  node [shape=rounded, style=filled, fillcolor=white, fontname=\"Arial\"];\n");
        sb.append("  edge [fontname=\"Arial\", fontsize=10];\n\n");

        // --- Core State Machine Rendering ---
        renderCore(result.getTransitions(), result.getStartStates(), result.getEndStates(), options, sb);

        sb.append("}\n");
        return sb.toString();
    }

    private void renderCore(List<Transition> transitions, Set<String> startStates, Set<String> endStates, ExportOptions options, StringBuilder sb) {
        // Choices
        Set<String> statesWithChoice = new LinkedHashSet<>();
        Map<String, String> choiceColorMap = new HashMap<>();
        int colorIndex = 0;

        for (Transition t : transitions) {
            if (t.getType() == TransitionType.CHOICE && t.getSourceStates() != null) {
                for (State source : t.getSourceStates()) {
                    String simplified = simplify(options.formatState(source));
                    statesWithChoice.add(simplified);
                    if (!choiceColorMap.containsKey(simplified)) {
                        choiceColorMap.put(simplified, CHOICE_COLORS[colorIndex % CHOICE_COLORS.length]);
                        colorIndex++;
                    }
                }
            }
        }

        boolean includeChoiceStates = options.isRenderChoicesAsDiamonds() && !statesWithChoice.isEmpty();
        if (includeChoiceStates) {
            for (String choice : statesWithChoice) {
                String color = choiceColorMap.get(choice);
                sb.append("  ").append(choice).append("_choice [shape=diamond, label=\"\", fillcolor=\"").append(color).append("\", width=0.2, height=0.2];\n");
                sb.append("  ").append(choice).append(" -> ").append(choice).append("_choice [style=dotted, color=\"").append(color).append("\"];\n");
            }
        }

        // Start/End nodes
        sb.append("  _start [shape=circle, label=\"\", fillcolor=black, width=0.1];\n");
        for (String start : startStates) {
            sb.append("  _start -> ").append(simplify(options.formatStateLabel(start))).append(";\n");
        }
        for (String end : endStates) {
            sb.append("  ").append(simplify(options.formatStateLabel(end))).append(" [fillcolor=lightgray];\n");
        }

        // Transitions
        for (Transition t : transitions) {
            if (t.getSourceStates() == null || t.getSourceStates().isEmpty())
                continue;

            boolean allowNoTarget = t.getType() == TransitionType.JOIN || t.getType() == TransitionType.FORK;
            List<State> targets = t.getTargetStates();

            if (targets == null || targets.isEmpty()) {
                if (!allowNoTarget) {
                    targets = t.getSourceStates(); // self-loop
                } else {
                    continue;
                }
            }

            for (State rawSourceState : t.getSourceStates()) {
                String rawSource = options.formatState(rawSourceState);
                String source = simplify(rawSource);

                for (State rawTargetState : targets) {
                    String rawTarget = options.formatState(rawTargetState);
                    String target = simplify(rawTarget);
                    String edgeSource = source;

                    if (t.getType() == TransitionType.CHOICE && statesWithChoice.contains(source)) {
                        edgeSource = includeChoiceStates ? source + "_choice" : source;
                    }

                    // Label: event [guard] / actions
                    StringBuilder label = new StringBuilder();
                    if (t.getEvent() != null) {
                        String eventStr = options.formatEvent(t.getEvent());
                        if (eventStr != null && !eventStr.isBlank()) {
                            label.append(escapeDotString(eventStr));
                        }
                    }

                    // Guards
                    if (t.getGuards() != null && !t.getGuards().isEmpty()) {
                        if (!label.isEmpty())
                            label.append(" ");
                        label.append("[");
                        for (int i = 0; i < t.getGuards().size(); i++) {
                            if (i > 0) label.append(" && ");
                            var guard = t.getGuards().get(i);
                            String guardText = options.isUseLambdaGuards() && guard.isLambda() ? "λ"
                                    : escapeDotString(guard.expression().replaceAll("[\\n\\r]+", " ").replaceAll("\\s{2,}", " ").trim());
                            label.append(guardText);
                        }
                        label.append("]");
                    }

                    // Actions
                    if (t.getActions() != null && !t.getActions().isEmpty()) {
                        StringBuilder actionBuilder = new StringBuilder();
                        for (int i = 0; i < t.getActions().size(); i++) {
                            if (i > 0)
                                actionBuilder.append(", ");
                            var action = t.getActions().get(i);
                            actionBuilder.append(options.isUseLambdaGuards() && action.isLambda() ? "λ" : escapeDotString(action.expression()));
                        }
                        if (!label.isEmpty())
                            label.append(" / ");
                        label.append(actionBuilder);
                    }

                    // Order
                    if ((t.getType() == TransitionType.CHOICE || t.getType() == TransitionType.JUNCTION) && t.getOrder() != null) {
                        if (!label.isEmpty())
                            label.append(" ");
                        label.append("(order=").append(t.getOrder()).append(")");
                    }

                    // Edge attributes
                    StringBuilder edgeAttrs = new StringBuilder();
                    if (!label.isEmpty()) {
                        appendEdgeAttr(edgeAttrs, "label", label.toString());
                    }

                    // Style/color by type
                    if (t.getType() != null) {
                        switch (t.getType()) {
                            case INTERNAL:
                                appendEdgeAttr(edgeAttrs, "style", "dashed");
                                appendEdgeAttr(edgeAttrs, "color", "gray");
                                break;
                            case LOCAL:
                                appendEdgeAttr(edgeAttrs, "style", "dotted");
                                appendEdgeAttr(edgeAttrs, "color", "darkgray");
                                break;
                            case CHOICE:
                                String choiceColor = choiceColorMap.getOrDefault(source, "blue");
                                appendEdgeAttr(edgeAttrs, "style", "solid");
                                appendEdgeAttr(edgeAttrs, "color", choiceColor);
                                break;
                            case JUNCTION:
                                appendEdgeAttr(edgeAttrs, "style", "dotted");
                                appendEdgeAttr(edgeAttrs, "color", "purple");
                                break;
                            case FORK:
                                appendEdgeAttr(edgeAttrs, "style", "bold");
                                appendEdgeAttr(edgeAttrs, "color", "green");
                                break;
                            case JOIN:
                                appendEdgeAttr(edgeAttrs, "style", "bold");
                                appendEdgeAttr(edgeAttrs, "color", "darkorange");
                                break;
                            case EXTERNAL:
                            default:
                                appendEdgeAttr(edgeAttrs, "style", "solid");
                                appendEdgeAttr(edgeAttrs, "color", "black");
                        }
                    } else {
                        appendEdgeAttr(edgeAttrs, "style", "solid");
                        appendEdgeAttr(edgeAttrs, "color", "black");
                    }

                    sb.append("  ").append(edgeSource).append(" -> ").append(target);
                    if (!edgeAttrs.isEmpty()) {
                        sb.append(" [").append(edgeAttrs).append("]");
                    }
                    sb.append(";\n");
                }
            }
        }
    }

    @Override
    public String export(String name,
                         List<Transition> transitions,
                         Set<String> startStates,
                         Set<String> endStates,
                         ExportOptions options) {
        StringBuilder sb = new StringBuilder();
        sb.append("digraph statemachine {\n");
        sb.append("  rankdir=LR;\n");
        sb.append("  node [shape=rounded, style=filled, fillcolor=white, fontname=\"Arial\"];\n");
        sb.append("  edge [fontname=\"Arial\", fontsize=10];\n\n");
        
        renderCore(transitions, startStates, endStates, options, sb);
        
        sb.append("}\n");
        return sb.toString();
    }

    @Override
    public String getFileExtension() {
        return ".dot";
    }

    private String escapeDotString(String s) {
        if (s == null)
            return "";
        return s.replace("\"", "\\\"");
    }

    private void appendEdgeAttr(StringBuilder sb, String key, String value) {
        if (!sb.isEmpty())
            sb.append(", ");
        sb.append(key).append("=\"").append(value).append("\"");
    }
}
