yao.scala

import java.io.{DataOutputStream, ByteArrayOutputStream}
import java.nio.ByteBuffer
import java.security.{KeyPairGenerator, SecureRandom}
import javax.crypto.{KeyGenerator, Cipher}
import javax.crypto.spec.IvParameterSpec

import akka.actor.{ ActorRef, ActorSystem, Props, Actor }

sealed trait Message
case class Start(visavis: ActorRef)
case class EncryptedDelta(n: Long)
case class Cryptos(p: Long, n: Seq[Long])
case class MoreOrEqual(name: String)
case class Less(name: String)

object Crypto {
  val iv = Array[Byte](0x0a, 0x01, 0x02, 0x03, 0x0b, 0x0c, 0x0d, 0x0a, 0x01, 0x02, 0x03, 0x04, 0x0b, 0x0c, 0x0b, 0x0a)
  val keyGenerator = KeyGenerator.getInstance("AES")
  keyGenerator.init(128)
  val key = keyGenerator.generateKey()

  def encrypt(n: Long): Long = {
    val cipher = Cipher.getInstance("AES/CTR/NoPadding")
    val ips = new IvParameterSpec(iv)
    cipher.init(Cipher.ENCRYPT_MODE, key, ips)
    bytesToLong(cipher.doFinal(longToBytes(n)))
  }

  def decrypt(n: Long): Long = {
    val cipher = Cipher.getInstance("AES/CTR/NoPadding")
    val ips = new IvParameterSpec(iv)
    cipher.init(Cipher.DECRYPT_MODE, key, ips)
    bytesToLong(cipher.doFinal(longToBytes(n)))
  }

  def bytesToLong(bytes: Array[Byte]) = {
    ByteBuffer.wrap(bytes).getLong
  }

  def longToBytes(n: Long) = {
    val baos = new ByteArrayOutputStream
    val dos = new DataOutputStream(baos)
    dos.writeLong(n)
    dos.close()
    baos.toByteArray
  }
}

object Prime {
  def isPrime(n: Long): Boolean = (n > 1) && (stream takeWhile { _ <= Math.sqrt(n) } forall { n % _ != 0 })
  val stream = Stream.cons[Long](2, Stream.from(3, 2).map(_.toLong) filter isPrime)

  def nextRandom: Long = {
    BigInt(7 * 8, 10, util.Random).toLong
  }
}

class Party(secret: Int, maxVariants: Int) extends Actor {
  val x = Prime.nextRandom
  println(s"${self.path.name}'s x: $x")
  val k = Crypto.encrypt(x)
  println(s"${self.path.name}'s k: $k")

  def receive = {
    case Start(visavis) => initiateExchange(visavis)
    case EncryptedDelta(delta) => {
      val name = self.path.name
      println(s"$name received encrypted delta value: $delta")
      val (p, variants) = proposedVariants(delta)
      println(s"$name calculated propsed variants of size ${variants.size} using prime $p")
      sender ! Cryptos(p, variants)
    }
    case Cryptos(p, variants) => sender ! decideIfBigger(p, variants)
    case MoreOrEqual(name) => println(s"$name is richer or same")
    case Less(name) => println(s"$name is poorer")
    case _ => println("WTF?!")
  }

  def decideIfBigger(p: Long, variants: Seq[Long]) = {
    val expected = x % p
    println(s"${self.path.name} expect $expected to be at position $secret")
    println(variants)
    if (variants(secret) == expected) {
      Less(self.path.name)
    } else {
      MoreOrEqual(self.path.name)
    }
  }

  def initiateExchange(visavis: ActorRef): Unit = {
    val selfName = self.path.name
    val visavisName = visavis.path.name
    println(s"$selfName was forced to start communication with $visavisName")
    visavis ! EncryptedDelta(k - secret)
  }

  def proposedVariants(delta: Long) = {
    val variants = for (u <- 0 to maxVariants - 1) yield Crypto.decrypt(delta + u)
    println(s"Variants: $variants")
    val (p, mods) = modVariants(variants)
    println(s"modVariants: $mods")
    val (left, right) = mods.splitAt(secret + 1)
    println(s"Splitted modVariants for $secret: $left, $right")
    val rightInc = for (i <- right) yield i + 1
    (p, left ++ rightInc)
  }

  final def modVariants(bare: Seq[Long], p: Long = Prime.nextRandom): (Long, Seq[Long]) = {
    println("calculating modVariants")
    val divided = for (i <- bare) yield i % p
    val diffs = for (a <- divided; b <- divided) yield math.abs(a - b)
    (p, divided)
  }

}

object Party {
  def props(secret: Int, maxVariants: Int) = Props(new Party(secret, maxVariants))
}

object YaoProtocol extends App {
  val system = ActorSystem("yao")

  val MAX = 2000

  val alice = system.actorOf(Party.props(args(0).toInt, MAX), "alice")
  val bob = system.actorOf(Party.props(args(1).toInt, MAX), "bob")
  bob ! Start(alice)

  def secretNumber(max: Int = 10) = util.Random.nextInt(max)
}