datatype instruction = NOOP | DO | DONT | MUL of (int * int)

fun numberAt (i : int) = fn (s : string) =>
  if i >= String.size s then NONE
  else SOME (Substring.takel Char.isDigit (Substring.extract (s, i, NONE)))

fun startsWith (needle : string, haystack : string, pos : int) =
  size haystack > pos + size needle andalso String.substring (haystack, pos, size needle) = needle

fun substrToInt (n : substring) = valOf (Int.fromString (Substring.string n))

val file = TextIO.openIn "day_3.txt"
val input = TextIO.inputAll file;
TextIO.closeIn file;

val instructions = List.tabulate ((size input), (fn i =>
  if startsWith ("do()", input, i) then DO
  else if startsWith ("don't()", input, i) then DONT
  else
    let
      val n1 = numberAt (i + 4) input
      val n2 = Option.mapPartial (fn n => numberAt (i + 5 + Substring.size n) input) n1
    in
      case (n1, n2) of
        (SOME(a), SOME(b)) =>
          if
            startsWith ("mul(", input, i) andalso
            String.sub (input, i + 4 + Substring.size a) = #"," andalso
            String.sub (input, i + 5 + Substring.size a + Substring.size b) = #")"
          then MUL (substrToInt a, substrToInt b) else NOOP
         | _ => NOOP
    end
))

val instructions = List.filter (fn (ins) => ins <> NOOP) instructions

fun calculate (instructions : instruction vector) =
  Vector.foldl (fn (ins, v) => case ins of MUL (a, b) => v + a * b | _ => v) 0 instructions

(* Part 1 *)

val mulOnlyResult = calculate (Vector.fromList instructions)

(* Part 2 *)

val filteredInstructions = Vector.mapi
  (fn (i, ins) =>
    let
      val prevInstructions = List.rev (List.take (instructions, i))
      val lastToggle = List.find (fn (ins) => ins = DO orelse ins = DONT) prevInstructions
      val shouldMultiply = case lastToggle of
        SOME (ins) => ins = DO
          | NONE => true
    in
      case ins of
        MUL (a, b) => if shouldMultiply then ins else NOOP
         | _ => ins
    end
  )
  (Vector.fromList instructions)

val allInsResult = calculate filteredInstructions