package com.sixtyfour.cbmnative.crossoptimizer.passes;

import com.sixtyfour.Logger;
import com.sixtyfour.cbmnative.crossoptimizer.common.OrderedPCode;
import com.sixtyfour.cbmnative.crossoptimizer.common.PCodeVisitor;
import com.sixtyfour.elements.commands.*;
import com.sixtyfour.parser.Line;

import java.util.List;
import java.util.SortedSet;
import java.util.TreeSet;

/**
 * This transformation will join lines that are not targeted by GOTOs
 *
 * The main benefit is that if you have a line with many commands and there is no loop in them, a constant propagator would be safe to apply.
 *
 * Code transforms from:
 * 10 PRINT "A"
 * 20 GOTO 50
 * 50 ....
 *
 * 150 GOTO 10
 *
 * To:
 *
 * 10 PRINT "A": GOTO 50
 * 50 ...
 * 150 GOTO 10
 *
 */
public class GenerateBasicBlocks {

    private static class Analysis {
        public SortedSet<Integer> rowsWithGotoTarget = new TreeSet<>();
        public SortedSet<Integer> rowsWithJumps = new TreeSet<>();

        public void addRow(int row) {
            rowsWithGotoTarget.add(row);
        }

        public void clear() {
            rowsWithGotoTarget.clear();
        }

        public void addStatement(Goto gotoStatement) {
            addRow(gotoStatement.getTargetLineNumber());
        }

        public void addStatement(Gosub gotoStatement) {
            addRow(gotoStatement.getTargetLineNumber());
        }

        public void addCommand(Command command) {
            if (command instanceof Gosub) {
                addStatement((Gosub) command);
            } else if (command instanceof Goto) {
                addStatement((Goto) command);
            } else {
                //command with no goto
                return;
            }
        }

        public void checkForJumps(int lineNumber, Command command) {
            if (isJumpCommand(command)) {
                rowsWithJumps.add(lineNumber);
            }
        }

        public static boolean isLineContainingJumps(Line line) {
            for (Command c : line.getCommands()) {
                if (isJumpCommand(c)) {
                    return true;
                }
            }
            return false;
        }

        private static boolean isJumpCommand(Command command) {
            if (command instanceof Gosub) {
                return true;
            }
            if (command instanceof Goto) {
                return true;
            }
            if (command instanceof Return) {
                return true;
            }
            return false;
        }

        public boolean isLineWithJumps(Line line) {
            return rowsWithJumps.contains(line.getNumber());
        }

        public boolean isLineTargetedByJump(Line line) {
            return rowsWithGotoTarget.contains(line.getNumber());
        }
    }

    private Analysis analysis = new Analysis();

    public boolean optimize(OrderedPCode orderedPCode) {
        boolean result = false;
        analyze(orderedPCode);
        List<Line> allLines = orderedPCode.getLines();
        for (int rowIndex = 0; rowIndex < allLines.size() - 1; rowIndex++) {
            Line currentRow = orderedPCode.getLineDirect(rowIndex);
            if (analysis.isLineWithJumps(currentRow)) {
                continue;
            }
            boolean found = true;

            while (found && (rowIndex + 1 < allLines.size())) {
                found = false;
                Line nextRow = orderedPCode.getLineDirect(rowIndex + 1);
                if (analysis.isLineTargetedByJump(nextRow)) {
                    continue;
                }
                result = true;
                found = true;
                joinTwoLines(orderedPCode, rowIndex);
                if (Analysis.isLineContainingJumps(currentRow))
                    break;
            }
        }
        return result;
    }

    private void joinTwoLines(OrderedPCode orderedPCode, int rowIndex) {
        assert orderedPCode.getLines().size() > rowIndex + 1;
        Line currentRow = orderedPCode.getLineDirect(rowIndex);
        Line nextRow = orderedPCode.getLineDirect(rowIndex + 1);
        currentRow.getCommands().addAll(nextRow.getCommands());
        currentRow.setLine(currentRow.getLine() + ":" + nextRow.getLine());
        Logger.log("Rows: " + currentRow.getNumber() + " and " + nextRow.getNumber()
                + " were merged as: '" + currentRow.getLine() + "'");
        orderedPCode.removeRow(nextRow.getNumber());
    }

    private void analyze(OrderedPCode orderedPCode) {
        PCodeVisitor visitor = new PCodeVisitor();
        visitor.accept(orderedPCode, (line, command, index) -> {
            analysis.addCommand(command);
            analysis.checkForJumps(line.getNumber(), command);
        });
    }
}


@@ -289,16 +289,44 @@
 .Line;%0A%0A
+import java.util.ArrayList;%0A
 import j
@@ -393,16 +393,52 @@
 TreeSet;
+%0Aimport java.util.stream.Collectors;
 %0A%0A/**%0A *
@@ -3070,24 +3070,78 @@
 getLines();%0A
+        List%3CInteger%3E rowsMerged = new ArrayList%3C%3E();%0A
         for 
@@ -3385,30 +3385,65 @@
     
-boolean found = true;%0A
+rowsMerged.clear();%0A            rowsMerged.add(rowIndex);
 %0A   
@@ -3460,18 +3460,8 @@
 hile
- (found &&
  (ro
@@ -3493,43 +3493,11 @@
 e())
-)
  %7B%0A
-                found = false;%0A
     
@@ -3643,32 +3643,29 @@
             
-continue
+break
 ;%0A          
@@ -3707,38 +3707,8 @@
 ue;%0A
-                found = true;%0A
     
@@ -3746,32 +3746,44 @@
 dPCode, rowIndex
+, rowsMerged
 );%0A             
@@ -3833,16 +3833,34 @@
 ntRow))%0A
+                %7B%0A
         
@@ -3882,33 +3882,351 @@
 ak;%0A            
-%7D
+    %7D%0A            %7D%0A            final String joinedRowsStr = rowsMerged.stream().map(value -%3E value.toString())%0A                    .collect(Collectors.joining(%22, %22));%0A            final String logMessage = %22Rows: %22 + joinedRowsStr + %22 were merged as: '%22 + currentRow.getLine() + %22'%22;%0A            Logger.log(logMessage);
 %0A        %7D%0A     
@@ -4319,16 +4319,42 @@
 rowIndex
+, List%3CInteger%3E rowsMerged
 ) %7B%0A    
@@ -4692,152 +4692,42 @@
     
-Logger.log(%22Rows: %22 + currentRow.getNumber() + %22 and %22 + nextRow.getNumber()%0A                + %22 were merged as: '%22 + currentRow.getLine() + %22'%22
+rowsMerged.add(nextRow.getNumber()
 );%0A 

Improve logging and make sure it merges a bit less aggressive