import qualified Data.Set as S
import Enigma
import Enigma.Internal
import EnigmaTestStatic
import Test.Hspec

-- Test helper
stepNTimes :: EnigmaConfig -> Int -> EnigmaConfig
stepNTimes cfg n =
  if n == 0 then cfg else stepNTimes (Enigma.Internal.step cfg) (n - 1)

main :: IO ()
main = hspec $ do
  describe "Enigma.Internal.index" $ do
    it "Gets the 0-25 index of a rotor top letter" $ do
      Enigma.Internal.index 'A' `shouldBe` (0 :: Int)
      Enigma.Internal.index 'Z' `shouldBe` (25 :: Int)

  describe "Enigma.Internal.revIndex" $ do
    it "Gets the top letter from an index num in 0-25" $ do
      Enigma.Internal.revIndex 0 `shouldBe` 'A'
      Enigma.Internal.revIndex 25 `shouldBe` 'Z'
      Enigma.Internal.revIndex (Enigma.Internal.index 'A') `shouldBe` 'A'

  describe "Enigma.Internal.mapRotor" $ do
    it "Maps a rotor input to an output for a given wiring" $ do
      -- Right to left mappings
      Enigma.Internal.mapRotor rotorIDWiring 'A' 0 `shouldBe` 0
      -- Identity map should work regardless of rotor position
      Enigma.Internal.mapRotor rotorIDWiring 'B' 0 `shouldBe` 0
      -- Test input at the last contact
      Enigma.Internal.mapRotor rotorIWiring 'A' 25 `shouldBe` 9
      -- Output contact - offet is negative and the output pos needs
      -- to "roll" over to a high index
      -- We have offset = 20, rhContact = 22, lhContact = 1, lh_pos = 7
      Enigma.Internal.mapRotor rotorIWiring 'U' 2 `shouldBe` 7
      Enigma.Internal.mapRotor rotorIIIWiring 'Z' 1 `shouldBe` 2
      Enigma.Internal.mapRotor rotorIWiring 'A' 0 `shouldBe` 4
      Enigma.Internal.mapRotor rotorIWiring 'B' 0 `shouldBe` 9

      -- Left to right mappings
      -- Input at last contact
      Enigma.Internal.mapRotor invRotorIWiring 'A' 25 `shouldBe` 9
      -- Offset is negative and rolls over to an index at the start of the rotor
      Enigma.Internal.mapRotor invRotorIWiring 'U' 12 `shouldBe` 11
      -- Rotor top is Z and input is 25 - causing overflow and underflow in
      -- the lhContact and rh output calculation
      Enigma.Internal.mapRotor invRotorIIWiring 'Z' 25 `shouldBe` 22
      Enigma.Internal.mapRotor invRotorIWiring 'A' 0 `shouldBe` 20
      Enigma.Internal.mapRotor invRotorIWiring 'F' 10 `shouldBe` 14

  describe "Enigma.Internal.mapReflector" $ do
    it "Maps reflector inputs to outputs" $ do
      Enigma.Internal.mapReflector rotorIDWiring 0 `shouldBe` 0
      -- Check that reflector is an involution
      Enigma.Internal.mapReflector refBWiring 4 `shouldBe` 16
      Enigma.Internal.mapReflector refBWiring 16 `shouldBe` 4

      -- Test the last position in the reflector
      Enigma.Internal.mapReflector refCWiring 11 `shouldBe` 25

      Enigma.Internal.mapReflector refCWiring 25 `shouldBe` 11
      Enigma.Internal.mapReflector refBWiring 0 `shouldBe` 24

  describe "Enigma.Internal.mapPlug" $ do
    it "Maps plugboard IO" $ do
      Enigma.Internal.mapPlug plugboardA 'Z' `shouldBe` 'A'
      Enigma.Internal.mapPlug plugboardB 'Z' `shouldBe` 'A'
      -- Plugboard mapping is an involution
      Enigma.Internal.mapPlug plugboardA (Enigma.Internal.mapPlug plugboardA 'X') `shouldBe` 'X'
      -- Character is missing from plugboard
      Enigma.Internal.mapPlug plugboardA 'Q' `shouldBe` 'Q'
      -- NB: Will need to change this test if representation of plugboard changes
      Enigma.Internal.mapPlug [] 'A' `shouldBe` 'A'

  describe "Enigma.Internal.cipherChar" $ do
    it "Enciphers a single character" $ do
      Enigma.Internal.cipherChar configA 'O' `shouldBe` 'D'
      Enigma.Internal.cipherChar configA 'Z' `shouldBe` 'H'
      Enigma.Internal.cipherChar configA 'Q' `shouldBe` 'V'
      Enigma.Internal.cipherChar configA 'V' `shouldBe` 'Q'
      Enigma.Internal.cipherChar configB 'R' `shouldBe` 'V'
      Enigma.Internal.cipherChar configB 'J' `shouldBe` 'Q'
      -- Check the case where signal passes through the plugboard twice
      Enigma.Internal.cipherChar configB 'X' `shouldBe` 'E'

      Enigma.Internal.cipherChar idConfig 'A' `shouldBe` 'A'
      Enigma.Internal.cipherChar configA 'G' `shouldBe` 'P'

  describe "Enigma.topLetters" $ do
    it "Extracts top letters from an EnigmaConfig" $ do
      Enigma.topLetters configA `shouldBe` "AAA"
      Enigma.topLetters configB `shouldBe` "ZQB"
      Enigma.topLetters stepCfgA `shouldBe` "KDO"
      Enigma.topLetters (stepNTimes stepCfgA 3) `shouldBe` Enigma.topLetters stepCfgA3

  describe "Enigma.Internal.stepRotors" $ do
    it "Steps a set of rotors" $ do
      Enigma.Internal.step stepCfgA `shouldBe` stepCfgA'
      stepNTimes stepCfgA 3 `shouldBe` stepCfgA3
      head (Enigma.topLetters (stepNTimes stepCfgB 30000)) `shouldBe` head (Enigma.topLetters stepCfgB)

  describe "Enigma.Internal.cipher" $ do
    it "Enciphers a string" $ do
      Enigma.cipher stepCfgA "OCAML" `shouldBe` "VOMUZ"
      Enigma.cipher stepCfgA "VOMUZ" `shouldBe` "OCAML"
      Enigma.cipher stepCfgA "HASKELL" `shouldBe` "ZLFIHSS"
      Enigma.cipher stepCfgA "ZLFIHSS" `shouldBe` "HASKELL"
      -- Long test cases verified against the Universal Enigma simulator
      -- http://people.physik.hu-berlin.de/~palloks/js/enigma/enigma-u_v25_en.html
      let plain = "OHMAMACANTHISREALLYBETHEENDTOBESTUCKINSIDEOFMOBILEWITHTHEMEMPHISBLUESAGAIN"
          ctext = "VGAHGCRJEZXTNQIXVACNAZMPYBZVJNYLIVAEWVNOMGQCZMQVWDCSYWRONWYEYSCCRFNPLEKILF"
      Enigma.cipher stepCfgA plain `shouldBe` ctext
      Enigma.cipher stepCfgA ctext `shouldBe` plain

      -- Tests incorporating the plugboard
      let plain2 = "COMEYOUMASTERSOFWARYOUTHATBUILDALLTHEGUNSYOUTHATBUILDTHEDEATHPLANES"
          ctext2 = "IHDVEWBUWODYNENLOUEZHHKKOLUBYOOKDFOPPKCWFZRANGWMNQAVLNUGISVXVDIMSPC"
      Enigma.cipher plugCfgA plain2 `shouldBe` ctext2

    describe "EnigmaInternal.freqs" $ do
      it "Returns a list of the number of times an item appears in a foldable" $ do
        S.fromList (Enigma.Internal.freqs "AABC") `shouldBe` S.fromList [2, 1, 1]
        Enigma.Internal.freqs "" `shouldBe` []
        Enigma.Internal.freqs "AAAAA" `shouldBe` [5]

    describe "Enigma.Internal.ic" $ do
      it "Calculates the index of coincidence from a list of frequencies" $ do
        Enigma.Internal.ic [1, 1] `shouldBe` 0
        Enigma.Internal.ic [2, 2] `shouldBe` 1 / 3
        Enigma.Internal.ic [2, 2, 2] `shouldBe` 1 / 5
        Enigma.Internal.ic (Enigma.Internal.freqs "AABBCC") `shouldBe` 1 / 5

    describe "Enigma.Internal.normalize" $ do
      it "Normalizes input text to [A..Z]" $ do
        Enigma.normalize "" `shouldBe` ""
        Enigma.normalize "ABCDE" `shouldBe` "ABCDE"
        Enigma.normalize ['a' .. 'z'] `shouldBe` ['A' .. 'Z']
        Enigma.normalize "é.~!@#$%^&*()" `shouldBe` ""
        Enigma.normalize ['0' .. '9'] `shouldBe` ""
        Enigma.normalize "This is a test string." `shouldBe` "THISISATESTSTRING"
