package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"os"
	"slices"
	"strconv"
	"unicode"

	"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) []postfixNode {
	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 (or an escape character)
			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_.
	*/
	i := 0
	for i < len(re_runes) {
		re_postfix = append(re_postfix, re_runes[i])
		if re_runes[i] == '[' && (i == 0 || re_runes[i-1] != '\\') { // We do not touch things inside brackets, unless they are escaped
			re_postfix[len(re_postfix)-1] = LBRACKET // Replace the '[' character with LBRACKET. This allows for easier parsing of all characters (including opening and closing brackets) within the character class
			invertMatch := false
			toAppend := make([]rune, 0)                      // Holds all the runes in the current character class
			if i < len(re_runes)-1 && re_runes[i+1] == '^' { // Inverting class - match everything NOT in brackets
				invertMatch = true
				i++
			}
			if i < len(re_runes)-1 && re_runes[i+1] == ']' { // Nothing inside brackets - panic.
				panic("Empty character class.")
			}
			for re_runes[i] != ']' {
				i++ // Skip all characters inside brackets
				// TODO: Check for escaped characters

				// Check ahead for character range
				if i < len(re_runes)-2 && re_runes[i+1] == '-' {
					rangeStart := re_runes[i]
					rangeEnd := re_runes[i+2]
					if int(rangeEnd) < int(rangeStart) {
						panic("Range is out of order.")
					}

					for i := rangeStart; i <= rangeEnd; i++ {
						toAppend = append(toAppend, i)
					}

					i += 2 // Skip start and hyphen (end will automatically be skipped on next iteration of loop)
					continue
				}
				toAppend = append(toAppend, re_runes[i])
			}
			// Replace the last character (which should have been ']', with RBRACKET
			toAppend[len(toAppend)-1] = RBRACKET
			if invertMatch {
				toAppend = setDifference(dotChars(), toAppend) // Take the inverse of the set by getting the difference between it and all dot characters
				toAppend = append(toAppend, RBRACKET)          // Since RBRACKET doesn't exist in dotChars, it wouldn't have been return in setDifference. We manually append it here.
			}
			re_postfix = append(re_postfix, toAppend...)
		}
		if re_runes[i] == '{' && (i > 0 && re_runes[i-1] != '\\') { // We don't touch things inside braces, either
			i++ // Skip opening brace
			for i < len(re_runes) && re_runes[i] != '}' {
				re_postfix = append(re_postfix, re_runes[i])
				i++
			}
			if i == len(re_runes) {
				panic("Invalid numeric specifier.")
			}
			re_postfix = append(re_postfix, re_runes[i]) // Append closing brace
		}
		if (re_runes[i] != '(' && re_runes[i] != '|' && re_runes[i] != '\\') || (i > 0 && re_runes[i-1] == '\\') { // Every character should be concatenated if it is escaped
			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_runes[i+1] != '{' {
					re_postfix = append(re_postfix, CONCAT)
				}
			}
		}
		i++
	}

	opStack := make([]rune, 0)         // Operator stack
	outQueue := make([]postfixNode, 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.
		5. If current character is '[', find all the characters until ']', then create a postfixNode containing all these contents. Add this node to outQueue.
		6. If current character is '{', find the appropriate numeric specifier (range start, range end). Apply the range to the postfixNode at the end of outQueue.
		*/
		c := re_postfix[i]
		if isAlphaNum(c) {
			outQueue = append(outQueue, newPostfixNode(c))
			continue
		}
		// Escape character
		if c == '\\' { // Escape character - invert special and non-special characters eg. \( is treated as a literal parentheses, \b is treated as word boundary
			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.")
			}
			i++
			outQueue = append(outQueue, newEscapedNode(re_postfix[i]))
			continue // Escaped character will automatically be skipped when loop variable increments
		}

		if c == '.' { // Dot metacharacter - represents 'any' character, but I am only adding Unicode 0020-007E
			outQueue = append(outQueue, newPostfixDotNode())
			continue
		}
		if c == '^' { // Start-of-string assertion
			outQueue = append(outQueue, newPostfixNode(c))
		}
		if c == '$' { // End-of-string assertion
			outQueue = append(outQueue, newPostfixNode(c))
		}

		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, newPostfixNode(to_append))
						topStack, _ = peek(opStack)
					}
					opStack = append(opStack, c)
				}
			}
		}
		if c == LBRACKET { // Used for character classes
			i++                      // Step forward so we can look at the character class
			chars := make([]rune, 0) // List of characters -  used only for character classes
			for i < len(re_postfix) {
				if re_postfix[i] == RBRACKET {
					break
				}
				chars = append(chars, re_postfix[i])
				i++
			}
			if i == len(re_postfix) { // We have reached the end of the string, so we didn't encounter a closing brakcet. Panic.
				panic("ERROR: Opening bracket without closing bracket.")
			}
			outQueue = append(outQueue, newPostfixNode(chars...))
			//	i++ // Step forward to skip closing bracket
			continue
		}
		if c == '{' {
			i++ // Skip opening brace
			// Three possibilities:
			// 1. Single number - {5}
			// 2. Range - {3,5}
			// 3. Start with no end, {3,}
			startRange := make([]rune, 0)
			startRangeNum := 0
			endRange := make([]rune, 0)
			endRangeNum := 0
			for i < len(re_postfix) && unicode.IsDigit(re_postfix[i]) {
				startRange = append(startRange, re_postfix[i])
				i++
			}
			if len(startRange) == 0 { // {} is not valid, neither is {,5}
				panic("ERROR: Invalid numeric specifier.")
			}
			if i == len(re_postfix) {
				panic("ERROR: Brace not closed.")
			}

			startRangeNum, err := strconv.Atoi(string(startRange))
			if err != nil {
				panic(err)
			}

			if re_postfix[i] == '}' { // Case 1 above
				endRangeNum = startRangeNum
			} else {
				if re_postfix[i] != ',' {
					panic("ERROR: Invalid numeric specifier.")
				}
				i++ // Skip comma
				for i < len(re_postfix) && unicode.IsDigit(re_postfix[i]) {
					endRange = append(endRange, re_postfix[i])
					i++
				}
				if i == len(re_postfix) {
					panic("ERROR: Brace not closed.")
				}
				if re_postfix[i] != '}' {
					panic("ERROR: Invalid numeric specifier.")
				}
				if len(endRange) == 0 { // Case 3 above
					endRangeNum = INFINITE_REPS
				} else { // Case 2 above
					var err error
					endRangeNum, err = strconv.Atoi(string(endRange))
					if err != nil {
						panic(err)
					}
				}
			}
			node, err := pop(&outQueue)
			if err != nil {
				panic("Numeric specifier with no content.")
			}
			node.startReps = startRangeNum
			node.endReps = endRangeNum
			outQueue = append(outQueue, node)
		}
		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, newPostfixNode(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, newPostfixNode(to_append))
	}

	return outQueue
}

// Thompson's algorithm. Constructs Finite-State Automaton from given string.
// Returns start state.
func thompson(re []postfixNode) *State {
	nfa := make([]*State, 0) // Stack of states
	for _, c := range re {
		if c.nodetype == CHARACTER || c.nodetype == ASSERTION {
			state := State{}
			state.transitions = make(map[int][]*State)
			if c.isDot {
				state.isDot = true
			}
			state.content = rune2Contents(c.contents)
			state.output = make([]*State, 0)
			state.output = append(state.output, &state)
			state.isEmpty = false
			if c.nodetype == ASSERTION {
				state.content = newContents(EPSILON) // Ideally, an assertion shouldn't have any content, since it doesn't say anything about the content of string
				state.isEmpty = true
				switch c.contents[0] {
				case '^':
					state.assert = SOS
				case '$':
					state.assert = EOS
				case 'b':
					state.assert = WBOUND
				case 'B':
					state.assert = NONWBOUND
				}
			}
			nfa = append(nfa, &state)
		}
		// Must be an operator if it isn't a character
		switch c.nodetype {
		case CONCATENATE:
			s2 := mustPop(&nfa)
			s1 := mustPop(&nfa)
			s1 = concatenate(s1, s2)
			nfa = append(nfa, s1)
		case KLEENE: // 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 PLUS: // a+ is equivalent to aa*
			s1 := mustPop(&nfa)
			s2 := kleene(*s1)
			s1 = concatenate(s1, s2)
			nfa = append(nfa, s1)
		case QUESTION: // ab? is equivalent to a(b|)
			s1 := mustPop(&nfa)
			s2 := question(s1)
			nfa = append(nfa, s2)
		case PIPE:
			s1 := mustPop(&nfa)
			s2 := mustPop(&nfa)
			s3 := alternate(s1, s2)
			nfa = append(nfa, s3)
		}
		if c.startReps != 1 || c.endReps != 1 { // Must have a numeric specifier attached to it
			if c.endReps != -1 && c.endReps < c.startReps {
				panic("ERROR: Numeric specifier - start greater than end.")
			}
			state := mustPop(&nfa)
			var stateToAdd *State = nil
			// Take advantage of the following facts:
			// a{5} == aaaaa
			// a{3,5} == aaaa?a?
			// a{5,} == aaaaa+
			// Nov. 3 2024 - I have two choices on how I want to implement numeric
			// specifiers.
			// a. Encode the logic while creating the states. I will have to create a function
			// that creates a deep-copy of a given state / NFA, so that I can concatenate them to
			// each other (concatenating them with the 'concatenate' method - which takes addresses - does
			// not work). Creating this function might be a lot of work.
			// b. Encode the logic while parsing the string (shunting-yard). If I can expand the numeric specifier
			// at this point, I can leave thompson untouched.
			for i := 0; i < c.startReps; i++ { // Case 1
				stateToAdd = concatenate(stateToAdd, cloneState(state))
			}
			if c.endReps == INFINITE_REPS { // Case 3
				s2 := kleene(*state)
				stateToAdd = concatenate(stateToAdd, s2)
			} else { // Case 2
				for i := c.startReps; i < c.endReps; i++ {
					stateToAdd = concatenate(stateToAdd, question(state))
				}
			}
			nfa = append(nfa, stateToAdd)
		}
	}
	if len(nfa) != 1 {
		panic("ERROR: Invalid Regex.")
	}

	verifyLastStates(nfa)

	return nfa[0]

}

func main() {
	invertFlag := flag.Bool("v", false, "Invert match.")
	// This flag has two 'modes':
	// 1. Without '-v': Prints only matches. Prints a newline after every match.
	// 2. With '-v': Substitutes all matches with empty string.
	onlyFlag := flag.Bool("o", false, "Print only colored content. Overrides -l.")
	lineFlag := flag.Bool("l", false, "Only print lines with a match (or with no matches, if -v is enabled")
	flag.Parse()

	// -l and -o are mutually exclusive: -o overrides -l
	if *onlyFlag {
		*lineFlag = false
	}

	// 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(flag.Args()) != 1 { // flag.Args() also strips out program name
		fmt.Println("ERROR: Missing cmdline args")
		os.Exit(22)
	}
	var re string
	re = flag.Args()[0]
	var test_str string
	var test_runes []rune // Rune-slice representation of test_str
	var err error
	// Create reader for stdin and writer for stdout // End index is one more than last index of match
	reader := bufio.NewReader(os.Stdin)
	out := bufio.NewWriter(os.Stdout)

	re_postfix := shuntingYard(re)
	startState := thompson(re_postfix)
	// Read every string from stdin until we encounter an error. If the error isn't EOF, panic.'
	for test_str, err = reader.ReadString('\n'); err == nil; test_str, err = reader.ReadString('\n') {
		test_runes = []rune(test_str)
		matchIndices := findAllMatches(startState, []rune(test_runes))
		// Decompose the array of matchIndex structs into a flat unique array of ints - if matchIndex is {4,7}, flat array will contain 4,5,6
		// This should make checking O(1) instead of O(n)
		indicesToPrint := new_uniq_arr[int]()
		for _, idx := range matchIndices {
			indicesToPrint.add(genRange(idx.startIdx, idx.endIdx)...)
		}
		// If we are inverting, then we should print the indices which _didn't_ match
		// in color.
		if *invertFlag {
			oldIndices := indicesToPrint.values()
			indicesToPrint = new_uniq_arr[int]()
			// Explanation:
			// Find all numbers from 0 to len(test_str) that are NOT in oldIndices.
			// These are the values we want to print, now that we have inverted the match.
			// Re-initialize indicesToPrint and add all of these values to it.
			indicesToPrint.add(setDifference(genRange(0, len(test_runes)), oldIndices)...)

		}
		// If lineFlag is enabled, we should only print something if:
		// 		a. We are not inverting, and have at least one match on the current line
		// 		OR
		// 		b. We are inverting, and have no matches at all on the current line.
		// This checks for the inverse, and continues if it is true.
		if *lineFlag {
			if !(*invertFlag) && len(matchIndices) == 0 || *invertFlag && len(matchIndices) > 0 {
				continue
			}
		}
		for i, c := range test_runes {
			if indicesToPrint.contains(i) {
				color.New(color.FgRed).Fprintf(out, "%c", c)
				// Newline after every match - only if -o is enabled and -v is disabled.
				if *onlyFlag && !(*invertFlag) {
					for _, idx := range matchIndices {
						if i+1 == idx.endIdx { // End index is one more than last index of match
							fmt.Fprintf(out, "\n")
							break
						}
					}
				}
			} else {
				if !(*onlyFlag) {
					fmt.Fprintf(out, "%c", c)
				}
			}
		}
		err = out.Flush()
		if err != nil {
			panic(err)
		}
	}
	if err != io.EOF {
		panic(err)
	}
}