# Install the required package, if it is not already installed.
import Pkg; Pkg.add("Random")

# We connect a package for generating random numbers and shuffling
using Random

"""
    struct ADFGVX

Defines the structure for storing ADFGVX cipher parameters:

- `polybius': symbol table (encrypted alphabet)
- `pdim`: Polybius table size (e.g. 6 for 6x6)
- `key`: the key of the transposition
- `keylen': key length
- `alphabet`: алфавит, используемый для кодирования (обычно "ADFGVX")
- `encode': encryption dictionary (character -> letter pair)
- `decode': decryption dictionary (letter pair -> symbol)
"""
struct ADFGVX
    polybius::Vector{Char}
    pdim::Int
    key::Vector{Char}
    keylen::Int
    alphabet::Vector{Char}
    encode::Dict{Char, Vector{Char}}
    decode::Dict{Vector{Char}, Char}
end

"""
    ADFGVX(s, k, alph = "ADFGVX")

The constructor creates an ADFGVX cipher object based on the data:
- `s`: a string of characters forming the Polybius table
- `k': the transposition key
- `alph`: алфавит шифра, по умолчанию "ADFGVX"
"""
function ADFGVX(s, k, alph = "ADFGVX")
    # Convert it to a vector of uppercase letters
    poly = collect(uppercase(s))
    pdim = isqrt(length(poly))  # calculating the dimension of the table (for a square one)
    
    al = collect(uppercase(alph)) # converting the alphabet
    
    # Creating an encryption dictionary: a pair of letters is assigned to each character.
    enco::Dict = Dict([(poly[(i - 1) * pdim + j] => [al[i], al[j]])
        for i in 1:pdim, j in 1:pdim])
    
    # Creating a dictionary for decoding (reverse)
    deco = Dict(last(p) => first(p) for p in enco)
    
    # Checking the correctness of the data
    @assert pdim^2 == length(poly) && pdim == length(al)
    
    # Returning the initialized object
    return ADFGVX(poly, pdim, collect(uppercase(k)), length(k), al, enco, deco)
end

"""
    encrypt(s::String, k::ADFGVX)

Encrypts the string `s` using an object `k` of type `ADFGVX'.
"""
function encrypt(s::String, k::ADFGVX)
    # Characters to be encrypted
    chars = reduce(vcat, [k.encode[c] for c in 
        filter(c -> c in k.polybius, collect(uppercase(s)))])

    # We distribute the symbols into columns according to the key
    colvecs = [lett => chars[i:k.keylen:length(chars)] for (i, lett) in enumerate(k.key)]
    
    # Arrange the columns alphabetically.
    sort!(colvecs, lt = (x, y) -> first(x) < first(y))
    
    # Combining the sorted columns into a row
    return String(mapreduce(p -> last(p), vcat, colvecs))
end

"""
    decrypt(s::String, k::ADFGVX)

Decrypts the string `s` encrypted with the ADFGVX cipher.
"""
function decrypt(s::String, k::ADFGVX)
    # Filtering the characters included in the cipher alphabet
    chars = filter(c -> c in k.alphabet, collect(uppercase(s)))
    sortedkey = sort(collect(k.key)) # sorting the key
    
    # Sorting order for columns
    order = [findfirst(c -> c == ch, k.key) for ch in sortedkey]
    originalorder = [findfirst(c -> c == ch, sortedkey) for ch in k.key]
    
    # Calculating column lengths (taking into account possible uneven distribution)
    a, b = divrem(length(chars), k.keylen)
    strides = [a + (b >= i ? 1 : 0) for i in order]
    starts = accumulate(+, strides[begin:end-1], init=1)
    pushfirst!(starts, 1)
    ends = [starts[i] + strides[i] - 1 for i in 1:k.keylen]
    
    # Restoring the columns to their original order
    cols = [chars[starts[i]:ends[i]] for i in originalorder]
    
    # Restoring rows from columns
    pairs, nrows = Char[], (length(chars) - 1) ÷ k.keylen + 1
    for i in 1:nrows, j in 1:k.keylen
        (i - 1) * k.keylen + j > length(chars) && break
        push!(pairs, cols[j][i])
    end
    
    # Converting pairs of characters back to the characters of the original alphabet
    return String([k.decode[[pairs[i], pairs[i + 1]]] for i in 1:2:length(pairs)-1])
end

const POLYBIUS = String(shuffle(collect("ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")))

# Download the English dictionary and select a random matching key of length 9
const KEY = read("unixdict.txt", String) |>
    v -> split(v, r"\s+") |>
    v -> filter(w -> (n = length(w); n == 9 && n == length(unique(collect(w)))), v) |>
    shuffle |> first |> uppercase

const SECRETS = ADFGVX(POLYBIUS, KEY)
message = "ATTACKAT1200AM"

# Data output to the screen
println("Polybius: $POLYBIUS, Key: $KEY\n")
println("Message: $message\n")

# Message encryption
encoded = encrypt(message, SECRETS)
println("Encoded: $encoded\n")

# Decrypting the message
decoded = decrypt(encoded, SECRETS)
println("Decoded: $decoded\n")