{-# LANGUAGE NamedFieldPuns #-}

module Enigma.Static where

import Data.Array (Array, accumArray, elems, listArray)
import Data.Char (chr, ord)
import Data.Maybe (fromJust)
import Enigma.Internal ( OrientedRotor(..), Rotor(..), Wiring )

-- NOTE: The order of this list is important. Do not modify
canonicalRotors :: [(String, (String, [Char]))]
canonicalRotors =
  [ ("I", ("EKMFLGDQVZNTOWYHXUSPAIBRCJ", "Q")),
    ("II", ("AJDKSIRUXBLHWTMCQGZNPYFVOE", "E")),
    ("III", ("BDFHJLCPRTXVZNYEIWGAKMUSQO", "V")),
    ("IV", ("ESOVPZJAYQUIRHXLNFTGKDCMWB", "J")),
    ("V", ("VZBRGITYUPSDNHLXAWMJQOFECK", "Z")),
    ("VI", ("JPGVOUMFYQBENHZRDKASXLICTW", "ZM")),
    ("VII", ("NZJHGRCXMYSWBOUFAIVLPEKQDT", "ZM")),
    ("VIII", ("FKQHTLXOCBJSPDZRAMEWNIUYGV", "ZM")),
    ("id", ("ABCDEFGHIJKLMNOPQRSTUVWXYZ", "")),
    ("refB", ("YRUHQSLDPXNGOKMIEBFZCWVJAT", "")),
    ("refC", ("FVPJIAOYEDRZXWGCTKUQSBNMHL", "")),
    ("beta", ("LEYJVCNIXWPBQMDRTAKZGFUHOS", "")),
    ("gamma", ("FSOKANUERHMBTIYCWLQPZXVGJD", ""))
  ]

invertWiring :: Wiring -> Wiring
invertWiring stringWiring =
  accumArray (\_ a -> a) '-' (0, 25) $ iv [] 0 (elems stringWiring)
  where
    iv :: [(Int, Char)] -> Int -> String -> [(Int, Char)]
    iv accum _ [] = accum
    iv accum ix (hd : tl) =
      iv ((ord hd - ord 'A', chr $ ix + ord 'A') : accum) (ix + 1) tl

-- | Returns of list of the five rotors used in the original
-- M3 Enigma Machine
m3RotorSet :: [Rotor]
m3RotorSet = buildRotorSet "VI"

buildRotorSet :: String -> [Rotor]
buildRotorSet stopAt = map (getRotor . fst) $ takeWhile (\(n, _) -> n /= stopAt) canonicalRotors

-- The pre-calculated reference IC of the Economist corpus
refIC :: Double
refIC = 0.0651437

-- | Get a rotor wiring by name i.e. wiring "I"
getWiring :: String -> Wiring
getWiring name = listArray (0, length w - 1) w
  where
    w = fst $ fromJust $ lookup name canonicalRotors

getTurnovers :: String -> [Char]
getTurnovers name = snd $ fromJust $ lookup name canonicalRotors

getRotor :: String -> Rotor
getRotor name = Rotor {rotId = name, wiring = w, invWiring = iw, turnovers = getTurnovers name}
  where
    w = getWiring name
    iw = invertWiring w

makeOriented :: String -> Char -> Char -> OrientedRotor
makeOriented name topLetter ringSetting =
  OrientedRotor {rotor = getRotor name, topLetter, ringSetting}
