;;;; Vigenere Cipher Encoder & Decoder program

(load "unit-test")
(defpackage :com.izak-halseide.vigenere-cipher
  (:use :common-lisp
        :com.izak-halseide.unit-test)
  (:export :encode
          :decode))
(in-package :com.izak-halseide.vigenere-cipher)

(defparameter *alphabet* "abcdefghijklmnopqrstuvwxyz")
(defparameter *alpha-size* (length *alphabet*))

;; Portable args fetching
(defun my-command-line ()
  "Return the command line args list (portable)"
  (or
     #+CLISP ext:*args*
     #+SBCL *posix-argv*
     #+LISPWORKS system:*line-arguments-list*
     #+CMU extensions:*command-line-words*
     nil))

(defun letter->num (letter)
  (if (stringp letter)
    (position (elt letter 0) *alphabet* :test 'char-equal)
    (position letter *alphabet* :test 'char-equal)))

(deftest test-letter->num ()
  (check
    (= (letter->num #\a) 0)
    (= (letter->num #\z) 25)
    (= (letter->num #\n) 13)))

(defun num->letter (num)
  (elt *alphabet* num))

(deftest test-num->letter ()
  (check
    (char= #\a (num->letter 0))
    (char= #\z (num->letter 25))
    (char= #\m (num->letter 12))))

(defun shift-letter (letter num)
  "Apply a Caesar cipher letter shift by num"
  ;; The LET and IF is necessary to make the result match the case of the input
  (let ((result (num->letter (mod (+ (letter->num letter) num) *alpha-size*))))
    (if (upper-case-p letter) (char-upcase result) result)))

(deftest test-shift-letter ()
  (check
    (char= #\b (shift-letter #\a 1))
    (char= #\a (shift-letter #\z 1))
    (char= #\z (shift-letter #\a -1))
    (char= #\a (shift-letter #\b -1))
    (char= #\P (shift-letter #\N 2))
    (char= #\P (shift-letter #\P 0))
    (char= #\s (shift-letter #\s 0))))

(defun encode (message key)
  (let ((output (make-array (length message) :element-type 'character :fill-pointer 0))) ; Create an appendable string
    (loop
      with key-len = (length key)
      with key-index = 0
      for letter across message do
        (if (alpha-char-p letter)
        ; then
        (let*
           ((key-value (letter->num (elt key key-index)))
            (new-letter (shift-letter letter key-value)))
          (vector-push new-letter output)
          ;; Is there a way to do this more concisely with LOOP?:
          (incf key-index)
          (if (>= key-index key-len) (setf key-index 0)))
        ; else
        (vector-push letter output)))
    output))

(deftest test-encode ()
  (check
    ;; Check example and that 'a' changes nothing
    (string= (encode "messages" "keys") "wiqkkkck")
    (string= (encode "Secret" "a") "Secret")
    ;; Only the message case should affect it
    (string= (encode "a" "a") "a")
    (string= (encode "A" "a") "A")
    (string= (encode "a" "A") "a")
    (string= (encode "A" "A") "A")))

(defun decode (message key)
  (let ((output (make-array (length message) :element-type 'character :fill-pointer 0))) ; Create an appendable string
    (loop
      with key-len = (length key)
      with key-index = 0
      for letter across message do
        (if (alpha-char-p letter)
        ; then
        (let*
           ((key-value (letter->num (elt key key-index)))
            (new-letter (shift-letter letter (- key-value))))
          (vector-push new-letter output)
          ;; Is there a way to do this more concisely with LOOP?:
          (incf key-index)
          (if (>= key-index key-len) (setf key-index 0)))
        ; else
        (vector-push letter output)))
    output))

(deftest test-decode ()
  (check
    (string= (decode "T" "a") "T")
    (string= (decode "Hflmo" "ab") "Hello"))
    (string= (decode "DARG" "bacon") "CAPS"))

(deftest test-encode-and-decode ()
  (check
    ;; Decoding an encoding should be the original message
    (string= (decode (encode "message" "ace") "ace") "message")))

(deftest test-everything ()
  (write-line "Testing...")
  (check
    (test-num->letter)
    (test-letter->num)
    (test-shift-letter)
    (test-encode)
    (test-decode)
    (test-encode-and-decode)))

(defun good-key-p (key-string)
  "A key must be only alphabetical letters"
  (= 0 (count-if-not #'alpha-char-p key-string)))

(defun split-args-for-cipher (args)
  "Takes the args and splits them into the expected message and key"
  (let ((message nil) (key (elt args (1- (length args)))))
    ;; If the first word is an argument (beginning with "-"), then we skip that
    ;; by subsequencing from a start of 1 instead of 0
    (setf message (subseq args (if (char= #\- (elt (first args) 0)) 1 0) (1- (length args))))
    (setf message (format nil "~{~A~^ ~}" message)) ; This joins the list into a string
    (list message key)))

(defun encode-args (args)
  (let ((arg-discovery (split-args-for-cipher args)))
    (format t "~a" (encode (elt arg-discovery 0) (elt arg-discovery 1)))))

(defun decode-args (args)
  (let ((arg-discovery (split-args-for-cipher args)))
    (format t "~a" (decode (elt arg-discovery 0) (elt arg-discovery 1)))))

(defun show-help ()
  (write-line "Showing help for vigenere-cipher...")
  (terpri)
  (write-line "Description:")
  (write-line "Encode or decode a message with the vigenere cipher, which means you use a word that is repeated across the length of the message and each letter shifts a letter of the plaintext by a certain amount. Special note: do not enter a message that begins with a dash (\"-\").")
  (terpri)
  (write-line "Usage: \"vigenere-cipher (-decode) message key\".")
  (terpri)
  (write-line "Options:")
  (write-line "Note that only the first option given will be used, and that each option can be abbreviated to one letter.")
  (write-line " -test   tests the program, and does not take any additional input.")
  (write-line " -help   shows this help.")
  (write-line " -?      is the same as \"-help\"")
  (write-line " -decode decodes a message instead of encoding a message.")
  (terpri))

(defun main ()
  (let ((args (my-command-line)))
    (cond
      ((= 0 (length args)) (show-help))
      ((string-equal "-h" (first args)) (show-help))
      ((string-equal "-?" (first args)) (show-help))
      ((string-equal "-help" (first args)) (show-help))
      ((string-equal "-t" (first args)) (test-everything))
      ((string-equal "-test" (first args)) (test-everything))
      ((string-equal "-d" (first args)) (decode-args args))
      ((string-equal "-decode" (first args)) (decode-args args))
      (t (encode-args args)))))

(main)