package ixee.cryptopals.utils.crypto

import javax.crypto.Cipher
import ixee.cryptopals.utils.ConversionUtils._
import ixee.cryptopals.utils.FunctionUtils._
import ixee.cryptopals.utils.CryptoUtils._
import ixee.cryptopals.utils.TupleUtils._
import ixee.cryptopals.utils.ByteUtils._

class EcbCipher(private[this] val cipher: Cipher, mode: Int) extends IxeeCipher {
  val blockSize = cipher.getBlockSize

  var leftover: Seq[Byte] = Seq()

  lazy val handleBlock: Seq[Byte] => Seq[Byte] =
    if (mode == Cipher.ENCRYPT_MODE) encBlock _
    else decBlock _

  def update(data: Seq[Byte]): Seq[Byte] = {
    val blocks = blockized(leftover ++ data).tap(updateLeftover)._1
    blocks.foldLeft(Seq[Byte]())(_ ++ handleBlock(_))
  }

  private def decBlock(data: Seq[Byte]): Seq[Byte] =
    cipher.update(data.toArray).toSeq

  private def encBlock(data: Seq[Byte]): Seq[Byte] =
    cipher.update(data.toArray).toSeq
  // wouldn't hurt to invalidate this object afterward, but meh
  // TODO: strip padding!
  // to do it right really requires writing decryption as its own part
  //   it's already obvious that's necessary, but to do padding stripping
  //   properly, the last block must be withheld until an end() call is made
  //   which is much different stateful behavior from encryption.
  //
  //   in cryptoUtils for now.
  def end(): Seq[Byte] =
    if (mode == Cipher.DECRYPT_MODE) Seq()
    else cipher.update(pkcs7pad(leftover, blockSize).toArray)

  def blockized(data: Seq[Byte]): (Seq[Seq[Byte]], Seq[Byte]) =
    groupBlocks <-: data.splitAt(data.length - (data.length % blockSize))

  def groupBlocks: Seq[Byte] => Seq[Seq[Byte]] = _.grouped(blockSize).toSeq

  def updateLeftover(pair: (Seq[Seq[Byte]], Seq[Byte])) =
    leftover = pair._2
}