package main

import (
	"flag"
	"log"
	"os"
	"regexp"
	"strings"
	"unicode"
)

const lettersInTheAlphabet = 'Z' - 'A' + 1

type keyLetter struct {
	upperCase string
	position  int32
}

func parseArguments() (string, string) {
	var inputKey string
	var inputFile string
	flag.StringVar(&inputKey, "k", "", "Specify key word.")
	flag.StringVar(&inputFile, "f", "", "Specify file to encrypt.")
	flag.Parse()
	return inputKey, inputFile
}

func readFile(relativePath string) string {
	content, err := os.ReadFile(relativePath)
	if err != nil {
		log.Fatalln(err)
	}
	return string(content)
}

func writeFile(relativePath string, message string) {
	file, err := os.Create(relativePath)
	if err != nil {
		log.Fatalln(err)
	}
	defer file.Close()
	file.WriteString(message)
}

func validateKeyWord(inputKey string) string {
	reg, err := regexp.Compile("[^a-zA-Z]+")
	if err != nil {
		log.Fatalln(err)
	}
	processedString := reg.ReplaceAllString(inputKey, "")
	return strings.ToUpper(processedString)
}

func initKeyWord(key string) []keyLetter {
	var keyWord []keyLetter
	for _, letter := range key {
		keyWord = append(keyWord, keyLetter{
			upperCase: string(letter),
			position:  letter - 'A',
		})
	}
	return keyWord
}

func encryptMessage(keyWord []keyLetter, message string) string {
	encrypted := ""
	index := 0
	for _, letter := range message {
		if !unicode.IsLetter(letter) {
			encrypted += string(letter)
		} else {
			keyWordIndex := index % len(keyWord)
			newLetter := letter + keyWord[keyWordIndex].position
			if (unicode.IsLower(letter) && newLetter > 'z') ||
				(unicode.IsUpper(letter) && newLetter > 'Z') {
				newLetter -= lettersInTheAlphabet
			}
			encrypted += string(newLetter)
			index += 1
		}
	}
	return encrypted
}

func main() {
	inputKey, inputFile := parseArguments()
	message := readFile(inputFile)
	keyWord := initKeyWord(validateKeyWord(inputKey))
	writeFile("output.txt", encryptMessage(keyWord, message))
}