package ru.yandex.atom.db.cassandra

import com.datastax.driver.core.{Row, Cluster}
import java.util.concurrent.Executors
import com.typesafe.config.Config
import scala.collection.JavaConversions._
import scala.concurrent.{ExecutionContext, Future, Promise}
import ru.yandex.atom.utils.concurrent.AsyncStream
import scala.util.Try
import ru.yandex.atom.db.cassandra.querybuilder.Statement

/**
 * @author avhaliullin
 */
trait CassandraComponentImpl extends CassandraComponent {

  class CassandraImpl(config: CassandraConfig) extends Cassandra {
    private val cluster = Cluster.builder().addContactPoints(config.nodes.map(_.url): _*)
      .withRetryPolicy(new CustomizableRetryPolicy(config.retryPolicy.readTORetries, config.retryPolicy.writeTORetries,
      config.retryPolicy.unavailableRetries))
      .build()
    private val session = cluster.connect(config.keyspace)
    private val executors = Executors.newFixedThreadPool(config.poolConfig.size)

    protected[this] implicit def executionContext = ExecutionContext.global

    def select(statement: Statement) = {
      val rsFuture = session.executeAsync(statement.statement)
      def makeListener: (Runnable, Future[AsyncStream[Seq[Row]]]) = {
        val p = Promise[AsyncStream[Seq[Row]]]()
        val listener = new Runnable {
          def run() {
            val res = Try(rsFuture.get()).map {
              rs =>
                val seq: Seq[Row] = for (i <- 0 until rs.getAvailableWithoutFetching) yield {
                  rs.one()
                }
                if (rs.isFullyFetched) {
                  AsyncStream(seq)
                } else {

                  AsyncStream(seq, {
                    val (nextListener, nextFuture) = makeListener
                    rsFuture.get().fetchMoreResults().addListener(nextListener, executors)
                    nextFuture
                  })
                }
            }
            p.tryComplete(res)
          }
        }
        (listener, p.future)
      }
      val (listener, future) = makeListener
      rsFuture.addListener(listener, executors)
      future
    }
  }

  object CassandraConfig {

    case class NodeConfig(url: String)

    case class ExecutorPoolConfig(size: Int)

    case class RetryPolicyConfig(readTORetries: Int, writeTORetries: Int, unavailableRetries: Int)

    implicit def apply(config: Config): CassandraConfig = {
      new CassandraConfig(
        nodes = config.getConfigList("nodes").map {
          nodeConfig =>
            new NodeConfig(nodeConfig.getString("url"))
        },
        keyspace = config.getString("keyspace"),
        poolConfig = new ExecutorPoolConfig(
          size = config.getInt("pool.size")
        ),
        retryPolicy = RetryPolicyConfig(
          readTORetries = config.getInt("retryPolicy.readTORetries"),
          writeTORetries = config.getInt("retryPolicy.writeTORetries"),
          unavailableRetries = config.getInt("retryPolicy.unavailableRetries")
        )
      )
    }
  }

  case class CassandraConfig(nodes: Seq[CassandraConfig.NodeConfig], keyspace: String,
                             poolConfig: CassandraConfig.ExecutorPoolConfig,
                             retryPolicy: CassandraConfig.RetryPolicyConfig) {
    assume(!nodes.isEmpty, "Cassandra should has nodes")
  }


}
