olympics.scala

import scala.collection.mutable._

// I want these std functions to be implicit:
implicit def tupled[a1, a2, b] = Function.tupled[a1, a2, b](_)
implicit def untupled[a1, a2, b] = Function.untupled[a1, a2, b](_)

class Connection

class ConnectionPool(size: Int) {
  def withConnection[A](f: Connection => A) = f(new Connection)
}

trait Query {
  def select: Seq[Int]
  def execute: Int
}

type QueryFactory = (Connection, String) => Query

case class SimpleQuery(connection: Connection, queryString: String) extends Query {
  def select = {
    Thread.sleep(1000)
    println("Selecting " + queryString + " on " + connection)
    Array(1, 2, 3)
  }

  def execute() = {
    Thread.sleep(1000)
    println("Executing " + queryString + " on " + connection)
    1
  }
  
  def use(thing: Int) = thing
}

trait QueryEvaluator {
  def select(queryString: String): Seq[Int]
  def execute(queryString: String): Int
  def transaction[A](f: QueryEvaluator => A): A
}

class InTransaction(Query: QueryFactory)(connection: Connection) extends QueryEvaluator {
  def select(queryString: String)  = Query(connection, queryString).select
  def execute(queryString: String) = Query(connection, queryString).execute

  def transaction[A](f: QueryEvaluator => A) = f(this)
}

class SimpleQueryEvaluator(Query: QueryFactory)(pool: ConnectionPool) {
  def select(queryString: String)  = pool.withConnection(Query(_, queryString).select)
  def execute(queryString: String) = pool.withConnection(Query(_, queryString).execute)

  def transaction[A](f: QueryEvaluator => A) = {
    pool.withConnection { c => f(new InTransaction(Query)(c)) }
  }
}

abstract class QueryProxy(protected var query: Query) extends Query {
  def select    = delegate("select")  { query.select }
  def execute() = delegate("execute") { query.execute }

  protected def delegate[A](name: String)(f: => A): A

  def reverse: QueryProxy = query match {
    case query: QueryProxy =>
      val reverse = query.reverse
      val inner = reverse.query
      this.query = inner
      reverse.query = this
      reverse
    case _ => this
  }
}

case class TimingOut(timeout: Int)(query: Query) extends QueryProxy(query) {
  def delegate[A](name: String)(f: => A) = {
    val result = f // Timeouts are more complicated in Java!
    println("Did not timeout! Yay fast database!")
    result
  }
}

class Stats {
  def measure[A](name: String)(f: => A) = {
    val start = System.currentTimeMillis
    val result = f
    val end = System.currentTimeMillis
    println("Measured " + name + " at " + (end - start) / 1000 + " seconds")
    result
  }
}

case class StatsCollecting(stats: Stats)(query: Query) extends QueryProxy(query) {
  def delegate[A](name: String)(f: => A) = stats.measure(name)(f)
}

case class Memoizing(Query: QueryFactory) extends QueryFactory {
  val memo = Map[(Connection, String), Query]()

  def apply(connection: Connection, queryString: String) = {
    memo.getOrElseUpdate((connection, queryString), Query(connection, queryString))
  }
}

case class Reversing(Query: QueryFactory) extends QueryFactory {
  def apply(connection: Connection, queryString: String) = Query(connection, queryString) match {
    case query: QueryProxy => query.reverse
    case query: Query => query
  }
}

val Query = Memoizing apply (SimpleQuery andThen TimingOut(1000) andThen StatsCollecting(new Stats))
for (Query <- List(Query, Reversing(Query))) {
  val queryEvaluator = new SimpleQueryEvaluator(Query)(new ConnectionPool(20))

  queryEvaluator.transaction { t =>
    t.select("SELECT ... FROM ... FOR UPDATE ...")
    t.execute("INSERT ...")
    t.execute("INSERT ...")
  }
  println()
}

// I don't really love the side-effecty reverse. But this actually has one fewer bug than the Ruby code
// Which I found thanks to the type-checker. This also took me about twice as long to write, though.