package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"unicode"
)

func main() {
	// Define the shift, hex, space, and grouping flags
	shift := flag.Int("s", 0, "Shift value for alphabet positions (can be negative)")
	useHex := flag.Bool("x", false, "Encode numbers in hexadecimal instead of decimal")
	encodeSpaces := flag.Bool("S", false, "Encodes spaces as '00'")
	groupSize := flag.Int("g", 2, "Group size for output digits (minimum value of 2, default 2)")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] <<< \"String\"\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "       %s [options] < file.txt\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "       echo \"string\" | %s [options]\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "\nDescription:\n")
		fmt.Fprintf(os.Stderr, "  A companion program to numberstation, intended to make easy work of making very basic ciphers for numbers station projects.\n")
		fmt.Fprintf(os.Stderr, "\nOptions:\n")
		flag.PrintDefaults() // Print the default flag help information
		fmt.Fprintf(os.Stderr, "\nExamples:\n")
		fmt.Fprintf(os.Stderr, "  %s -s 2 -g 5 <<< \"This is a test.\"\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "  %s -S -s 16 -x -g 4 < file.txt | numberstation -r\n", os.Args[0])
	}

	flag.Parse()

	// Check for invalid group size
	if *groupSize < 2 {
		fmt.Fprintln(os.Stderr, "Error: Minimum group size is 2. Exiting.")
		os.Exit(1)
	}

	// Check if input is being piped or passed as an argument
	stat, _ := os.Stdin.Stat()
	if (stat.Mode() & os.ModeCharDevice) != 0 {
		// No input, display help text
		flag.Usage()
		os.Exit(1)
	}

	// Read input from stdin
	inputBytes, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		fmt.Println("Error reading input:", err)
		os.Exit(1)
	}
	input := strings.TrimSpace(string(inputBytes))

	// Convert input to uppercase and prepare the result slice
	input = strings.ToUpper(input)
	var result []string
	var currentNumber string

	for _, char := range input {
		if unicode.IsLetter(char) {
			// If there was a number being processed, add it to the result first
			if currentNumber != "" {
				appendNumber(&result, currentNumber, *useHex, *groupSize)
				currentNumber = ""
			}
			// Get the position of the character in the alphabet (A = 1, B = 2, ..., Z = 26)
			position := int(char - 'A' + 1)
			// Apply the shift, wrapping around the alphabet
			shiftedPosition := (position + *shift - 1) % 26
			if shiftedPosition < 0 {
				shiftedPosition += 26
			}
			shiftedPosition += 1

			// Append the shifted position to the result slice
			result = append(result, fmt.Sprintf("%02d", shiftedPosition))
		} else if unicode.IsDigit(char) {
			// Collect digits to process as a whole number
			currentNumber += string(char)
		} else if char == ' ' && *encodeSpaces {
			// Encode space as '00' if the -S flag is set
			if currentNumber != "" {
				appendNumber(&result, currentNumber, *useHex, *groupSize)
				currentNumber = ""
			}
			result = append(result, "00")
		} else {
			// If there was a number being processed and now a non-digit/non-letter appears, append it
			if currentNumber != "" {
				appendNumber(&result, currentNumber, *useHex, *groupSize)
				currentNumber = ""
			}
		}
	}

	// If there's any remaining number at the end of the input, append it
	if currentNumber != "" {
		appendNumber(&result, currentNumber, *useHex, *groupSize)
	}

	// Format output with proper grouping
	finalOutput := groupOutput(result, *groupSize)
	fmt.Println(finalOutput)
}

// appendNumber appends a formatted number to the result slice
func appendNumber(result *[]string, number string, useHex bool, groupSize int) {
	var formatted string
	num, err := strconv.Atoi(number)
	if err != nil {
		return
	}

	// Set maximum values based on group size and mode
	var maxValue int
	if groupSize == 2 || groupSize == 4 {
		if useHex {
			maxValue = 255 // 0xFF
		} else {
			maxValue = 99 // 0d99
		}
	} else {
		maxDigits := groupSize - 2 // for "0x" or "0d" prefix
		maxValue = intMaxValue(useHex, maxDigits)
	}

	if num > maxValue {
		fmt.Fprintf(os.Stderr, "Warning: Number %d is too large for group size %d. Using maximum value %d.\n", num, groupSize, maxValue)
		num = maxValue
	}

	if useHex {
		formatted = fmt.Sprintf("0x%X", num)
	} else {
		formatted = fmt.Sprintf("0d%d", num)
	}

	*result = append(*result, formatted)
}

// intMaxValue calculates the maximum value that fits within the specified group size
func intMaxValue(useHex bool, maxDigits int) int {
	if useHex {
		return (1 << uint(maxDigits*4)) - 1 // Maximum hexadecimal value fitting in maxDigits
	}
	return intPow(10, maxDigits) - 1 // Maximum decimal value fitting in maxDigits
}

// intPow calculates integer power
func intPow(base, exp int) int {
	result := 1
	for exp > 0 {
		result *= base
		exp--
	}
	return result
}

// groupOutput groups the output into specified sizes, ensuring numbers are separate
func groupOutput(result []string, groupSize int) string {
	var output []string
	var currentGroup string

	for _, item := range result {
		if strings.HasPrefix(item, "0x") || strings.HasPrefix(item, "0d") {
			// If the current group is not empty, add it to the output
			if currentGroup != "" {
				output = append(output, currentGroup)
				currentGroup = ""
			}
			// Numbers are placed in their own group
			output = append(output, item)
		} else {
			currentGroup += item
			if len(currentGroup) >= groupSize {
				output = append(output, currentGroup[:groupSize])
				currentGroup = currentGroup[groupSize:]
			}
		}
	}

	// Append any remaining characters in the current group
	if currentGroup != "" {
		output = append(output, currentGroup)
	}

	return strings.Join(output, " ")
}