# 如果尚未安装，请安装所需的软件包。
import Pkg; Pkg.add("Random")

# 我们连接一个包，用于生成随机数和洗牌
using Random

"""
    struct ADFGVX

定义用于存储ADFGVX密码的参数的结构:

-'polybius'：符号表（加密字母表）
-'pdim':Polybius表尺寸(例如6x6)
-'键'：换位键
-'keylen'：键长度
- `alphabet`: алфавит, используемый для кодирования (обычно "ADFGVX")
-'encode'：加密字典（字符->字母对）
-'decode'：解密字典（字母对->符号）
"""
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")

构造函数根据数据创建ADFGVX密码对象:
-'s'：形成Polybius表的一串字符
-'k'：换位键
- `alph`: алфавит шифра, по умолчанию "ADFGVX"
"""
function ADFGVX(s, k, alph = "ADFGVX")
    # 将其转换为大写字母的向量
    poly = collect(uppercase(s))
    pdim = isqrt(length(poly))  # 计算表的维度（对于正方形）
    
    al = collect(uppercase(alph)) # 转换字母表
    
    # 创建加密字典：为每个字符分配一对字母。
    enco::Dict = Dict([(poly[(i - 1) * pdim + j] => [al[i], al[j]])
        for i in 1:pdim, j in 1:pdim])
    
    # 创建用于解码的字典（反向）
    deco = Dict(last(p) => first(p) for p in enco)
    
    # 检查数据的正确性
    @assert pdim^2 == length(poly) && pdim == length(al)
    
    # 返回初始化的对象
    return ADFGVX(poly, pdim, collect(uppercase(k)), length(k), al, enco, deco)
end

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

使用类型为`ADFGVX`的对象`k`加密字符串`s'。
"""
function encrypt(s::String, k::ADFGVX)
    # 要加密的字符
    chars = reduce(vcat, [k.encode[c] for c in 
        filter(c -> c in k.polybius, collect(uppercase(s)))])

    # 我们根据键将符号分配到列中
    colvecs = [lett => chars[i:k.keylen:length(chars)] for (i, lett) in enumerate(k.key)]
    
    # 按字母顺序排列列。
    sort!(colvecs, lt = (x, y) -> first(x) < first(y))
    
    # 将排序的列组合成一行
    return String(mapreduce(p -> last(p), vcat, colvecs))
end

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

解密使用ADFGVX密码加密的字符串's'。
"""
function decrypt(s::String, k::ADFGVX)
    # 过滤密码字母表中包含的字符
    chars = filter(c -> c in k.alphabet, collect(uppercase(s)))
    sortedkey = sort(collect(k.key)) # 排序密钥
    
    # 列的排序顺序
    order = [findfirst(c -> c == ch, k.key) for ch in sortedkey]
    originalorder = [findfirst(c -> c == ch, sortedkey) for ch in k.key]
    
    # 计算列长度（考虑到可能的不均匀分布）
    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]
    
    # 将列还原为其原始顺序
    cols = [chars[starts[i]:ends[i]] for i in originalorder]
    
    # 从列还原行
    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
    
    # 将字符对转换回原始字母表的字符
    return String([k.decode[[pairs[i], pairs[i + 1]]] for i in 1:2:length(pairs)-1])
end

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

# 下载英语词典并选择长度为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"

# 数据输出到屏幕
println("Polybius: $POLYBIUS, Key: $KEY\n")
println("Message: $message\n")

# 消息加密
encoded = encrypt(message, SECRETS)
println("Encoded: $encoded\n")

# 解密消息
decoded = decrypt(encoded, SECRETS)
println("Decoded: $decoded\n")