package test;

import org.antlr.v4.runtime.CharStream;
import org.antlr.v4.runtime.CharStreams;
import org.antlr.v4.runtime.CommonTokenStream;

import java.util.Random;

/**
 * Pointer-invariant test suite for the reference-based AST implementation
 * (parent / firstChild / tail / prevSibling / nextSibling).
 */
public final class AstOpsTest
{
   public static void main(String[] args) throws Exception
   {
      System.out.println("=== AST: ===");
      AnnotatedAst program = parse("x = 10 + 5 * 2;");
      printTree(program, 0);
      verifyPointers(program);

      AnnotatedAst assign = (AnnotatedAst) program.getFirstChild(); 
      AnnotatedAst xNode  = (AnnotatedAst) assign.getFirstChild();  
      AnnotatedAst expr   = (AnnotatedAst) xNode.getNextSibling(); 

      // TEST 1: null-safety
      System.out.println("\n=== TEST 1: null safety ===");
      assign.addChild(null);
      try
      {
         assign.graftAt(null, 0);
         assign.graft(null);
      }
      catch (NullPointerException e) 
      {
         System.out.println("NPE expected");
      }
      
      verifyPointers(program);

      // TEST 2: graftAt(-1) append
      System.out.println("\n=== TEST 2: graftAt(-1) should append node A as the last child ===");
      AnnotatedAst A = mk("A", 99);
      assign.graftAt(A, -1);
      verifyPointers(program);
      assertOrder(assign, new String[]{"x", "+", "A"}, "after graftAt(-1)");
      printTree(program, 0);

      // TEST 3: graftAt(very large) => append (fallback)
      System.out.println("\n=== TEST 3: graftAt(100) append fallback, node B must be the last child ===");
      AnnotatedAst B = mk("B", 99);
      assign.graftAt(B, 100);
      verifyPointers(program);
      assertOrder(assign, new String[]{"x", "+", "A", "B"}, "after graftAt(100)");
      printTree(program, 0);

      // TEST 4: move(null, 0) => move within same parent (B to front)
      System.out.println("\n=== TEST 4: move(null,0) should move node B to index 0 within the same parent ===");
      B.move(null, 0);
      verifyPointers(program);
      assertOrder(assign, new String[]{"B", "x", "+", "A"}, "after move(null,0)");
      printTree(program, 0);

      // TEST 5: move within same parent to end
      System.out.println("\n=== TEST 5: move(null,-1) to end. Node B must be the last child ===");
      B.move(null, -1);
      verifyPointers(program);
      assertOrder(assign, new String[]{"x", "+", "A", "B"}, "after move(null,-1)");
      printTree(program, 0);

      // TEST 6: move to a different parent: move 'A' under '+'
      System.out.println("\n=== TEST 6: move node 'A' from '=' to '+', '+' children should be: 10, *, A  ===");
      A.move(expr, -1);
      verifyPointers(program);
      assertOrder(assign, new String[]{"x", "+", "B"}, "after moving A out of '='");
      assertOrder(expr, new String[]{"10", "*", "A"}, "after moving A under '+'");
      printTree(program, 0);

      // TEST 7: remove middle child ('*' from '+')
      System.out.println("\n=== TEST 7: remove node '*' from '+' children (expected '+' children: 10, A) ===");
      AnnotatedAst star = (AnnotatedAst) expr.getFirstChild().getNextSibling(); // '*'
      star.remove();
      verifyPointers(program);
      assertOrder(expr, new String[]{"10", "A"}, "after removing '*'");
      printTree(program, 0);

      // TEST 8: setNextSibling(null) cuts chain (cut 'B' from '=')
      System.out.println("\n=== TEST 8: setNextSibling(null) on '+' (cut chain after '+', '=' children should be: x, +) ===");
      AnnotatedAst plus = (AnnotatedAst) xNode.getNextSibling(); // '+'
      plus.setNextSibling(null);
      verifyPointers(program);
      assertOrder(assign, new String[]{"x", "+"}, "after cutting chain at '+'");
      printTree(program, 0);


      // TEST 9: setFirstChild(null) removes all children (+ becomes leaf)
      System.out.println("\n=== TEST 9: setFirstChild(null) on '+' (remove all children, '+' becomes leaf) ===");
      plus.setFirstChild(null);
      verifyPointers(program);
      assertOrder(plus, new String[]{}, "after clearing children of '+'");
      printTree(program, 0);

      // TEST 10: remove root (no parent) is no-op
      System.out.println("\n=== TEST 10: remove root node 'PROGRAM' (should be no-op since it has no parent)===");
      program.remove();
      verifyPointers(program);
      printTree(program, 0);

      // TEST 11: randomized fuzz test
      System.out.println("\n=== TEST 11: fuzz test (random graft/move/remove operations, pointer invariants must always hold) ===");
      fuzzOnce(2000, 123);

      System.out.println("ALL TESTS PASSED :-D");
   }

   private static AnnotatedAst parse(String input)
   {
      CharStream cs = CharStreams.fromString(input);
      ExprLexer lexer = new ExprLexer(cs);
      CommonTokenStream tokens = new CommonTokenStream(lexer);
      ExprParser parser = new ExprParser(tokens);

      ExprParser.ProgramContext ctx = parser.program();
      return (AnnotatedAst) ctx.node;
   }

   private static AnnotatedAst mk(String text, int type)
   {
      AnnotatedAst n = new AnnotatedAst();
      n.setText(text);
      n.setType(type);
      return n;
   }

   private static void verifyPointers(AnnotatedAst node)
   {
      if (node == null)
      {
         return;
      }

      // verify children chain
      AnnotatedAst child = (AnnotatedAst) node.getFirstChild();
      AnnotatedAst last = null;

      while (child != null)
      {
         // parent pointer must match
         if (child.getParent() != node)
         {
            throw new RuntimeException("Parent mismatch at '" + child.getText()
                  + "': expected '" + node.getText() + "'");
         }

         // prevSibling chain must be consistent
         if (child.getPrevSibling() != last)
         {
            throw new RuntimeException("PrevSibling mismatch at '" + child.getText()
                  + "': expected '" + (last == null ? "null" : last.getText()) + "'");
         }

         last = child;
         child = (AnnotatedAst) child.getNextSibling();
      }

      // tail must match last (if tail is part of your implementation)
      AnnotatedAst tail = (AnnotatedAst) node.getTail();
      if (last != null && tail != last)
      {
         throw new RuntimeException("Tail mismatch at '" + node.getText()
               + "': expected '" + last.getText()
               + "', got '" + (tail == null ? "null" : tail.getText()) + "'");
      }
      if (last == null && tail != null)
      {
         throw new RuntimeException("Tail should be null at '" + node.getText() + "'");
      }

      // recurse
      child = (AnnotatedAst) node.getFirstChild();
      while (child != null)
      {
         verifyPointers(child);
         child = (AnnotatedAst) child.getNextSibling();
      }
   }

   private static void assertOrder(AnnotatedAst parent, String[] expected, String label)
   {
      int idx = 0;
      AnnotatedAst c = (AnnotatedAst) parent.getFirstChild();
      while (c != null)
      {
         if (idx >= expected.length)
         {
            throw new RuntimeException("Too many children " + label + " under '" + parent.getText()
                  + "'. Unexpected: '" + c.getText() + "'");
         }
         String exp = expected[idx];
         if (!exp.equals(c.getText()))
         {
            throw new RuntimeException("Child order mismatch " + label + " under '" + parent.getText()
                  + "': expected[" + idx + "]='" + exp + "', got '" + c.getText() + "'");
         }
         idx++;
         c = (AnnotatedAst) c.getNextSibling();
      }
      if (idx != expected.length)
      {
         throw new RuntimeException("Too few children " + label + " under '" + parent.getText()
               + "': expected " + expected.length + ", got " + idx);
      }
   }

   private static void printTree(AnnotatedAst node, int indent)
   {
      if (node == null) return;

      for (int i = 0; i < indent; i++) System.out.print("  ");
      System.out.println(node.getText() + " (Type: " + node.getType() + ")");

      AnnotatedAst child = (AnnotatedAst) node.getFirstChild();
      while (child != null)
      {
         printTree(child, indent + 1);
         child = (AnnotatedAst) child.getNextSibling();
      }
   }

   private static void fuzzOnce(int ops, int seed) throws Exception
   {
      Random rnd = new Random(seed);

      AnnotatedAst program = parse("x = 1 + 2 * 3;");
      verifyPointers(program);

      AnnotatedAst assign = (AnnotatedAst) program.getFirstChild();
      AnnotatedAst plus = (AnnotatedAst) ((AnnotatedAst) assign.getFirstChild()).getNextSibling();

      // keep some free-floating nodes to insert/move
      AnnotatedAst[] pool = new AnnotatedAst[50];
      for (int i = 0; i < pool.length; i++)
      {
         pool[i] = mk("N" + i, 777);
      }

      for (int i = 0; i < ops; i++)
      {
         int action = rnd.nextInt(6);
         switch (action)
         {
            case 0: // graft random pool node under '='
            {
               AnnotatedAst n = pool[rnd.nextInt(pool.length)];
               if (n.getParent() == null)
               {
                  assign.graftAt(n, rnd.nextBoolean() ? -1 : 0);
               }
               break;
            }
            case 1: // graft random pool node under '+'
            {
               AnnotatedAst n = pool[rnd.nextInt(pool.length)];
               if (n.getParent() == null)
               {
                  plus.graftAt(n, rnd.nextInt(5) - 1); // -1..3
               }
               break;
            }
            case 2: // move a random child within '=' (if any)
            {
               AnnotatedAst child = randomChild(assign, rnd);
               if (child != null)
               {
                  int idx = rnd.nextInt(4) - 1; // -1..2
                  child.move(null, idx);
               }
               break;
            }
            case 3: // move a random child from '=' to '+'
            {
               AnnotatedAst child = randomChild(assign, rnd);
               if (child != null)
               {
                  child.move(plus, rnd.nextInt(4) - 1);
               }
               break;
            }
            case 4: // remove a random child from '+'
            {
               AnnotatedAst child = randomChild(plus, rnd);
               if (child != null)
               {
                  child.remove();
               }
               break;
            }
            case 5: // cut chain at a random child in '='
            {
               AnnotatedAst child = randomChild(assign, rnd);
               if (child != null)
               {
                  child.setNextSibling(null);
               }
               break;
            }
            default:
               break;
         }

         // invariants must always hold
         verifyPointers(program);
      }
   }

   private static AnnotatedAst randomChild(AnnotatedAst parent, Random rnd)
   {
      int n = parent.getNumImmediateChildren();
      if (n == 0) return null;
      int idx = rnd.nextInt(n);
      return (AnnotatedAst) parent.getChildAt(idx);
   }
}