package ai.diffy.proxy

import java.io.File
import java.util.zip.ZipFile

import ai.diffy.analysis.{DifferenceAnalyzer, InMemoryDifferenceCollector, JoinedDifferences}
import ai.diffy.lifter.{MapLifterPool, Message, ThriftLifter}
import ai.diffy.scrooge.ZippedFileImporter
import com.twitter.finagle.thrift.{ClientId, ThriftClientRequest}
import com.twitter.finagle.tracing.NullTracer
import com.twitter.finagle.{Resolver, Thrift, ThriftMux}
import com.twitter.util.{Future, Try}

import scala.collection.JavaConversions._

case class ThriftDifferenceProxy (
    settings: Settings,
    collector: InMemoryDifferenceCollector,
    joinedDifferences: JoinedDifferences,
    analyzer: DifferenceAnalyzer)
  extends DifferenceProxy
{
  override type Req = ThriftClientRequest
  override type Rep = Array[Byte]
  override type Srv = ThriftService

  private[this] lazy val clientId = new ClientId(settings.clientId)

  override val proxy = super.proxy

  private[this] val zipfile = new ZipFile(new File(settings.pathToThriftJar))

  private[this] val importer =
    ZippedFileImporter(Seq(zipfile))

  private[this] val filenames =
    zipfile.entries.toSeq collect {
      case zipEntry if !zipEntry.isDirectory && zipEntry.getName.endsWith(".thrift") =>
        zipEntry.getName
    }

  val lifter =
    MapLifterPool(
      ThriftLifter.fromImporter(
        importer,
        filenames,
        settings.serviceClass
      )
    )

  override def serviceFactory(serverset: String, label: String) = {
    val client = if (settings.enableThriftMux) {
      ThriftMux.client
        .withClientId(clientId)
        .newClient(serverset, label).toService
    } else {
      val config = if(settings.useFramedThriftTransport) {
        Thrift.client
      } else {
        Thrift.client.withBufferedTransport
      }

      config
        .withNoAttemptTTwitterUpgrade
        .withTracer(NullTracer)
        .withClientId(clientId)
        .newClient(serverset, label)
        .toService
    }

    ThriftService(client, Resolver.eval(serverset))
  }

  override lazy val server = {
    if (settings.enableThriftMux) {
      ThriftMux.serve(
        settings.servicePort,
        proxy map { req: Array[Byte] => new ThriftClientRequest(req, false) }
      )
    } else {
      val config = if(settings.useFramedThriftTransport) {
        Thrift.server
      } else {
        Thrift.server.withBufferedTransport()
      }
      config.withTracer(NullTracer).serve(
        settings.servicePort,
        proxy map { req: Array[Byte] => new ThriftClientRequest(req, false) }
      )
    }
  }

  override def liftRequest(req: ThriftClientRequest): Future[Message] = lifter(req.message)
  override def liftResponse(rep: Try[Array[Byte]]): Future[Message] =
    Future.const(rep) flatMap { lifter(_) }
}