package gaf3.core.gateway.filter.factory

import org.reactivestreams.Publisher
import org.slf4j.LoggerFactory
import org.springframework.beans.factory.annotation.Value
import org.springframework.cloud.gateway.filter.GatewayFilter
import org.springframework.cloud.gateway.filter.GatewayFilterChain
import org.springframework.cloud.gateway.filter.NettyWriteResponseFilter
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils.ORIGINAL_RESPONSE_CONTENT_TYPE_ATTR
import org.springframework.cloud.gateway.support.ShortcutConfigurable.ShortcutType
import org.springframework.core.Ordered
import org.springframework.core.io.buffer.DataBuffer
import org.springframework.http.HttpHeaders
import org.springframework.http.MediaType
import org.springframework.http.server.reactive.ServerHttpResponseDecorator
import org.springframework.web.server.ServerWebExchange
import reactor.core.publisher.Flux
import reactor.core.publisher.Mono
import java.nio.charset.Charset


class PeekResponseBodyGatewayFilterFactory
    : AbstractGatewayFilterFactory<PeekResponseBodyGatewayFilterFactory.Config>(Config::class.java) {

    @Value("\${gateway.response.body.peek-size:2000}")
    var peekSize: Int = 2000

    override fun shortcutFieldOrder(): List<String> {
        return arrayListOf(MEDIA_TYPES_KEY)
    }

    override fun shortcutType(): ShortcutType {
        return ShortcutType.GATHER_LIST
    }

    override fun apply(factoryConfig: Config): GatewayFilter {
        return PeekResponseGatewayFilter(factoryConfig)
    }

    class Config { //Put the configuration properties for your filter here
        var mediaTypes: Array<String>? = null
            get() = if(field?.size ?: 0 > 0)
                field
            else
                arrayOf(MediaType.APPLICATION_JSON_VALUE,
                        MediaType.APPLICATION_XML_VALUE,
                        MediaType.TEXT_HTML_VALUE,
                        MediaType.TEXT_MARKDOWN_VALUE,
                        MediaType.TEXT_PLAIN_VALUE,
                        MediaType.TEXT_XML_VALUE)
    }

    inner class PeekResponseGatewayFilter(private val factoryConfig: Config): GatewayFilter, Ordered {
        override fun filter(exchange: ServerWebExchange, chain: GatewayFilterChain): Mono<Void> {

            return chain.filter(exchange.mutate()
                    .request { it.header(HttpHeaders.ACCEPT_ENCODING, "identity").build() }
                    .response(PeekedServerHttpResponse(exchange, factoryConfig)).build())
                    .then(Mono.fromRunnable {
                        log.debug("content-type: {}", exchange.response.headers.contentType)
                        val snapshot = exchange.attributes[ORIGINAL_RESPONSE_BODY_SNAPSHOT] as? String
                        log.debug("peek response body snapshot:\n{}", snapshot)
                    })
        }

        override fun getOrder(): Int {
            return NettyWriteResponseFilter.WRITE_RESPONSE_FILTER_ORDER - 1
        }
    }

    inner class PeekedServerHttpResponse(private val exchange: ServerWebExchange, private val factoryConfig: Config) : ServerHttpResponseDecorator(exchange.response) {

        override fun writeWith(body: Publisher<out DataBuffer>): Mono<Void> {
            val flux = Flux.from(body)
                    .replay()
                    .autoConnect()

            flux.take(1).subscribe { buffer: DataBuffer ->
                val contentType: MediaType? = exchange.response.headers.contentType ?:
                    exchange.getAttribute(ORIGINAL_RESPONSE_CONTENT_TYPE_ATTR)
                val supported = factoryConfig.mediaTypes?.any { mediaType ->
                    runCatching {
                        MediaType.parseMediaType(mediaType).isCompatibleWith(contentType)
                    }.getOrElse { err ->
                        log.warn("Invalid Media Type: $mediaType", err)
                        false
                    }
                } ?: false
                if(supported) {
                    val snapshot = buffer.toString(Charset.defaultCharset()).take(peekSize)
                    exchange.attributes[ORIGINAL_RESPONSE_BODY_SNAPSHOT] = snapshot
                }
            }
            return delegate.writeWith(flux)
        }

        override fun writeAndFlushWith(
                body: Publisher<out Publisher<out DataBuffer>>): Mono<Void> {
            return writeWith(Flux.from(body).flatMapSequential { p: Publisher<out DataBuffer>? -> p })
        }
    }

    companion object {
        const val MEDIA_TYPES_KEY = "mediaTypes"
        const val ORIGINAL_RESPONSE_BODY_SNAPSHOT = "ORIGINAL_RESPONSE_BODY_SNAPSHOT"

        internal val log = LoggerFactory.getLogger(PeekResponseBodyGatewayFilterFactory::class.java)
    }
}