package main

import (
	"fmt"
	"os"
	"slices"

	"github.com/fatih/color"
)

const CONCAT rune = '~'

func isOperator(c rune) bool {
	if c == '+' || c == '?' || c == '*' || c == '|' || c == CONCAT {
		return true
	}
	return false
}

/* priority returns the priority of the given operator */
func priority(op rune) int {
	precedence := []rune{'|', CONCAT, '+', '*', '?'}
	return slices.Index(precedence, op)
}

/*
The Shunting-Yard algorithm is used to convert the given infix (regeular) expression to postfix.
The primary benefit of this is getting rid of parentheses.
It also inserts explicit concatenation operators to make parsing easier in Thompson's algorithm.
See: https://blog.cernera.me/converting-regular-expressions-to-postfix-notation-with-the-shunting-yard-algorithm/
*/
func shuntingYard(re string) string {
	re_postfix := make([]rune, 0)
	re_runes := []rune(re) // Convert the string to a slice of runes to allow iteration through it
	/* 	Add concatenation operators.
	Only add a concatenation operator between two characters if both the following conditions are met:
		1. 	The first character isn't an opening parantheses or alteration operator.
			a. This makes sense, because these operators can't be _concatenated_ with anything else.
		2. The second character isn't a 'closing operator' - one that applies to something before it
			a. Again, these operators can'be concatenated _to_. They can, however, be concatenated _from_.
	*/
	for i := 0; i < len(re_runes); i++ {
		re_postfix = append(re_postfix, re_runes[i])
		if re_runes[i] != '(' && re_runes[i] != '|' {
			if i < len(re_runes)-1 {
				if re_runes[i+1] != '|' && re_runes[i+1] != '*' && re_runes[i+1] != '+' && re_runes[i+1] != '?' && re_runes[i+1] != ')' {
					re_postfix = append(re_postfix, CONCAT)
				}
			}
		}
	}

	opStack := make([]rune, 0)  // Operator stack
	outQueue := make([]rune, 0) // Output queue

	// Actual algorithm
	for i := 0; i < len(re_postfix); i++ {
		/* Two cases:
		1. Current character is alphanumeric - send to output queue
		2. Current character is operator - do the following:
			a. If current character has greater priority than top of opStack, push to opStack.
			b. If not, keep popping from opStack (and appending to outQueue) until:
				i. opStack is empty, OR
				ii. current character has greater priority than top of opStack
		3. If current character is '(', push to opStack
		4. If current character is ')', pop from opStack (and append to outQueue) until '(' is found. Discard parantheses.
		*/
		c := re_postfix[i]
		if isAlphaNum(c) {
			outQueue = append(outQueue, c)
			continue
		}
		// Escape character - NOT IMPLEMENTED YET - DO NOT USE
		// if c == '\\' { // Escape character - next character is treated as alphanum
		// 	if i == len(re_postfix)-1 { // End of string - panic, because backslash is an escape character (something needs to come after it)
		// 		panic("ERROR: Backslash with no escape character.")
		// 	}
		// 	outQueue = append(outQueue, re_postfix[i+1])
		// }

		if isOperator(c) {
			if len(opStack) == 0 {
				opStack = append(opStack, c)
			} else {
				topStack, err := peek(opStack)
				if err != nil {
					panic("ERROR: Operator without operand.")
				}
				if priority(c) > priority(topStack) { // 2a
					opStack = append(opStack, c)
				} else {
					for priority(c) <= priority(topStack) { // 2b
						to_append := mustPop(&opStack)
						outQueue = append(outQueue, to_append)
						topStack, _ = peek(opStack)
					}
					opStack = append(opStack, c)
				}
			}
		}
		if c == '(' {
			opStack = append(opStack, c)
		}
		if c == ')' {
			// Keep popping from opStack until we encounter an opening parantheses. Panic if we reach the end of the stack.
			for val, err := peek(opStack); val != '('; val, err = peek(opStack) {
				if err != nil {
					panic("ERROR: Imbalanced parantheses.")
				}
				to_append := mustPop(&opStack)
				outQueue = append(outQueue, to_append)
			}
			_ = mustPop(&opStack) // Get rid of opening parantheses
		}
	}

	// Pop all remaining operators (and append to outQueue)
	for len(opStack) > 0 {
		to_append := mustPop(&opStack)
		outQueue = append(outQueue, to_append)
	}

	return string(outQueue)
}

// Thompson's algorithm. Constructs Finite-State Automaton from given string.
// Returns start state.
func thompson(re string) *State {
	nfa := make([]*State, 0) // Stack of states
	for _, c := range re {
		if isAlphaNum(c) {
			state := State{}
			state.transitions = make(map[int][]*State)
			state.content = int(c)
			state.output = make([]*State, 0)
			state.output = append(state.output, &state)
			state.isEmpty = false
			nfa = append(nfa, &state)
		}
		// Must be an operator if it isn't alphanumeric
		switch c {
		case CONCAT:
			s2 := mustPop(&nfa)
			s1 := mustPop(&nfa)
			s1 = concatenate(s1, s2)
			nfa = append(nfa, s1)
		case '*': // Create a 0-state, concat the popped state after it, concat the 0-state after the popped state
			s1 := mustPop(&nfa)
			stateToAdd := kleene(*s1)
			nfa = append(nfa, stateToAdd)
		case '+': // a+ is equivalent to aa*
			s1 := mustPop(&nfa)
			s2 := kleene(*s1)
			s1 = concatenate(s1, s2)
			nfa = append(nfa, s1)
		case '?': // ab? is equivalent to a(b|)
			s1 := mustPop(&nfa)
			s2 := &State{}
			s2.transitions = make(map[int][]*State)
			s2.content = EPSILON
			s2.output = append(s2.output, s2)
			s2.isEmpty = true
			s3 := alternate(s1, s2)
			nfa = append(nfa, s3)
		case '|':
			s1 := mustPop(&nfa)
			s2 := mustPop(&nfa)
			s3 := alternate(s1, s2)
			nfa = append(nfa, s3)
		}
	}
	if len(nfa) != 1 {
		panic("ERROR: Invalid Regex.")
	}

	verifyLastStates(nfa)

	return nfa[0]

}

func main() {
	// Process:
	// 1. Convert regex into postfix notation (Shunting-Yard algorithm)
	// 		a. Add explicit concatenation operators to facilitate this
	// 2. Build NFA from postfix representation (Thompson's algorithm)
	// 3. Run the string against the NFA
	if len(os.Args) < 3 {
		fmt.Println("ERROR: Missing cmdline args")
		os.Exit(22)
	}
	var re string
	re = os.Args[1]
	re_postfix := shuntingYard(re)
	// fmt.Println(re_postfix)
	startState := thompson(re_postfix)
	matchIndices := findAllMatches(startState, os.Args[2])
	inColor := false
	if len(matchIndices) > 0 {
		for i, c := range os.Args[2] {
			for _, indices := range matchIndices {
				if i >= indices.startIdx && i < indices.endIdx {
					color.New(color.FgRed).Printf("%c", c)
					inColor = true
					break
				}
			}
			if inColor == false {
				fmt.Printf("%c", c)
			}
			inColor = false
		}
		fmt.Printf("\n")
	} else {
		fmt.Println(os.Args[2])
	}
}