abstract class Term
case class Var(s: String) extends Term
case class Fun(s: String, ts: List[Term]) extends Term


abstract class Form
object True extends Form
object False extends Form
case class Pred(s: String, ts: List[Term]) extends Form
case class Imp(f1: Form, f2: Form) extends Form
case class Says(p: String, f: Form) extends Form
case class And(f1: Form, f2: Form) extends Form
case class Or(f1: Form, f2: Form) extends Form
case class Controls(p: String, f: Form) extends Form

case class Judgement(Gamma: List[Form], F: Form) {
  def lhs = Gamma
  def rhs = F
}

val Alice = "Alice"
val Bob = "Bob"
val Send = Pred("Send", Nil)

val Gamma = 
  List( Imp(Says(Bob, Send), Send),
        Says(Bob, Imp(Says(Alice, Send), Send)),
        Says(Alice, Send) )

val goal = Judgement(Gamma, Send)

def sc () = { println ("Yes!") }

class Main {

def prove(j: Judgement, sc: () => Unit) : Unit = j match {
  case Judgement(lhs, rhs) => 
    { if (lhs.exists(f => f == rhs))  sc () 
      else prove1(lhs, rhs, sc) 
    }
}

def partitions [A] (l: List[A]): List[(A, List[A])]  = 
  l.map (s => (s, l - s))

def prove1(lhs: List[Form], rhs: Form, sc: () => Unit) : Unit = 
  rhs match {
    case Imp(f1, f2) => prove(Judgement(f1::lhs, f2), sc)
    /*case Says(p, f) => prove(Judgement(lhs, f), sc)*/
    /*case Controls(p, f) => prove(Judgement(lhs, f), sc)*/
    case Or(f1, f2) => { prove(Judgement(lhs, f1), sc);
                         prove(Judgement(lhs, f2), sc) }
    case And(f1, f2) => prove(Judgement(lhs, f1), 
                            () => prove(Judgement(lhs, f2), sc))
    case _ => { for ((f, lhs_rest) <- partitions(lhs))
                  prove2(f, lhs_rest, rhs, sc) }
  }

def prove2(f: Form, lhs_rest: List[Form], rhs: Form, sc: () => Unit) : Unit = 
  f match {
    case Imp(f1, f2) => 
      prove(Judgement(lhs_rest, f1), 
            () => prove(Judgement(f2::lhs_rest, rhs), sc))
    case Says(p, Imp(f1, f2)) => 
      prove(Judgement(lhs_rest, Says(p, f1)), 
            () => prove(Judgement(Says(p, f2)::lhs_rest, rhs), sc)) 
    case Controls(p, f) =>
      prove(Judgement(lhs_rest, Says(p, f)),
            () => prove(Judgement(f::lhs_rest, rhs), sc))
    case _ => ()
  }

}
  
val main = new Main
val Foo = Pred("Foo", Nil)

main.prove (Judgement (Gamma, And(Foo, Send)), sc)
main.prove (Judgement (Nil, Foo), sc)
main.prove (Judgement (Nil, Imp(Send, Send)), sc)
main.prove (Judgement (Gamma, Send), sc)
main.prove (Judgement (Gamma, Foo), sc)

val F1 = Imp(Says(Bob, Send), Send)
val F2 = Says(Bob, Imp(Says(Alice, Send), Send))
val F3 = Says(Alice, Send)

main.prove (Judgement (Nil, Imp(F1, Imp(F2, Imp(F3, Send)))), sc)

val Server = "Server"

def Sends(p: String, q: String, f: Form) : Form =
  Imp(Says(p, f), Says(q, f))

def Enc(f: Form, k: Form) : Form = Imp(k, f)

def Connect(p: String, q: String) : Form =
  Pred("Connect", List(Var(p), Var(q)))

val Msg = Pred("Msg", Nil)
val Kas = Pred("Kas", Nil)
val Kbs = Pred("Kbs", Nil)
val Kab = Pred("Kab", Nil)

val Gamma_big = 
  List( Says(Alice, Kas),
        Says(Bob, Kbs),
        Says(Alice, Msg),
        Says(Alice, Connect(Alice, Bob)),
        Sends(Alice, Server, Connect(Alice, Bob)),
        Says(Server, Imp(Connect(Alice, Bob), Enc(Kab, Kas))),
        Says(Server, Imp(Connect(Alice, Bob), Enc(Enc(Kab, Kbs), Kas))),
        Sends(Server, Alice, Enc(Kab, Kas)),
        Sends(Server, Alice, Enc(Enc(Kab, Kbs), Kas)),
        Sends(Alice, Bob, Enc(Kab, Kbs)),
        Sends(Alice, Bob, Enc(Msg, Kab))
      )


main.prove (Judgement(Gamma_big, Says(Bob, Msg)), sc)