package main

import (
	"bufio"
	"flag"
	"fmt"
	"os"
	"strings"
	"unicode"
)

// Function to encode characters into their respective number pairs
func encodeChar(char rune) int {
	// Define the alphabet and numbers with a leading space
	alphabet := " ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"

	// Find the index of the character
	index := strings.IndexRune(alphabet, char)
	if index == -1 {
		return -1 // Return -1 if character is not found
	}
	return index
}

// Function to apply shift and wrap around
func applyShift(value, shift int) int {
	shifted := (value + shift) % 37
	if shifted < 0 {
		shifted += 37
	}
	return shifted
}

// Function to format the output based on group size and line size
func formatOutput(encoded []int, groupSize int, lineSize int) string {
	var output strings.Builder
	for _, num := range encoded {
		output.WriteString(fmt.Sprintf("%02d", num))
	}

	finalOutput := output.String()
	var groupedOutput strings.Builder
	counter := 0
	groupCounter := 0

	// Check if a group size is provided
	if groupSize > 0 {
		// Apply grouping and line size limits
		for i := 0; i < len(finalOutput); i += groupSize {
			if i+groupSize < len(finalOutput) {
				groupedOutput.WriteString(finalOutput[i:i+groupSize] + " ")
			} else {
				groupedOutput.WriteString(finalOutput[i:])
			}
			counter++
			groupCounter++

			// Check if we reached the line size
			if lineSize > 0 && groupCounter >= lineSize && i+groupSize < len(finalOutput) {
				groupedOutput.WriteString("\n")
				groupCounter = 0
			}
		}
	} else {
		// No group size specified, handle line size by number of characters
		for i := 0; i < len(finalOutput); i++ {
			groupedOutput.WriteByte(finalOutput[i])
			counter++

			// Check if we reached the line size
			if lineSize > 0 && counter >= lineSize && i+1 < len(finalOutput) {
				groupedOutput.WriteString("\n")
				counter = 0
			}
		}
	}

	return strings.TrimSpace(groupedOutput.String())
}

func main() {
	// Define command-line flags
	shift := flag.Int("s", 0, "Shift amount for the cipher -- '-s 2' will encode A as 03 instead of 01")
	groupSize := flag.Int("g", 0, "Group size for output formatting -- '-g 5' will output numbers in groups of 5")
	lineSize := flag.Int("l", 0, "Line size by number of groups -- '-l 4' will split to a new line after every 4th group")

	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 4 <<< \"This is a test\"\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "  %s -s 16 -g 5 -l 4 < file.txt | numberstation -r\n", os.Args[0])
	}

	flag.Parse()

	// Check if there is input from stdin
	info, _ := os.Stdin.Stat()
	if (info.Mode() & os.ModeCharDevice) != 0 {
		flag.Usage()
		os.Exit(1)
	}

	// Read input from stdin
	reader := bufio.NewReader(os.Stdin)
	input, _ := reader.ReadString('\n')
	input = strings.TrimSpace(input)

	var encoded []int

	// Encode the input string
	for _, char := range input {
		if !unicode.IsLetter(char) && !unicode.IsDigit(char) && char != ' ' {
			continue // Skip characters not in the defined alphabet and numbers
		}
		encodedValue := encodeChar(unicode.ToUpper(char))
		if encodedValue == -1 {
			continue // Skip characters not found in the alphabet string
		}
		shiftedValue := applyShift(encodedValue, *shift)
		encoded = append(encoded, shiftedValue)
	}

	// Format the output
	formattedOutput := formatOutput(encoded, *groupSize, *lineSize)
	fmt.Println(formattedOutput)
}