package main

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
)

func indentMixCo(s string) (string, error) {
	pairs := map[string]string{
		"FOR": "NEXT",
		"IF":  "ENDIF",
	}
	level := 0
	lines := strings.Split(s, "\n")
	length, _ := strconv.Atoi(lines[0])
	spaces := strings.Trim(lines[1], "\n")
	lines = lines[2:]
	var stack []string
	var removed []string
	indented := false
	for i := 0; i < length; i++ {
		lines[i] = strings.TrimLeft(lines[i], "·»\t")
		for key, val := range pairs {
			rOpen, _ := regexp.Compile(fmt.Sprintf(`^%v($|\s)`, key))
			if rOpen.MatchString(lines[i]) {
				lines[i] = strings.Repeat(spaces, level) + lines[i]
				level++
				stack = append(stack, key)
				indented = true
			}
			rClose, _ := regexp.Compile(fmt.Sprintf(`^%v($|\s)`, val))
			if rClose.MatchString(lines[i]) {
				level--
				lines[i] = strings.Repeat(spaces, level) + lines[i]
				removed, stack = stack[len(stack)-1:], stack[:len(stack)-1]
				indented = true
				if removed[0] != key {
					return strings.Join(lines, "\n"), fmt.Errorf("error: %v found on line %v but was looking for %v", val, i+1, pairs[removed[0]])
				}
			}
		}
		if !indented {
			lines[i] = strings.Repeat(spaces, level) + lines[i]
		} else {
			indented = false
		}
	}
	return strings.Join(lines, "\n"), nil
}

func main() {
	input := `12
····
VAR I
·FOR I=1 TO 31
»»»»IF !(I MOD 3) THEN
··PRINT "FIZZ"
··»»ENDIF
»»»»····IF !(I MOD 5) THEN
»»»»··PRINT "BUZZ"
··»»»»»»ENDIF
»»»»IF (I MOD 3) && (I MOD 5) THEN
······PRINT "FIZZBUZZ"
··»»ENDIF
»»»»·NEXT`

	properFormat, err := indentMixCo(input)
	if err != nil {
		fmt.Println(err)
		fmt.Println(properFormat)
	} else {
		fmt.Println(properFormat)
	}
}
