package gaf3.core.gateway.filter.factory

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.factory.AbstractGatewayFilterFactory
import org.springframework.cloud.gateway.support.ShortcutConfigurable.ShortcutType
import org.springframework.core.Ordered
import org.springframework.core.io.buffer.DataBuffer
import org.springframework.http.MediaType
import org.springframework.http.server.reactive.ServerHttpRequestDecorator
import org.springframework.web.server.ServerWebExchange
import reactor.core.publisher.Flux
import reactor.core.publisher.Mono
import java.nio.charset.Charset


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

    @Value("\${gateway.request.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 PeekRequestGatewayFilter(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)
    }

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

            return chain.filter(exchange.mutate()
                    .request(PeekedServerHttpRequest(exchange, factoryConfig))
                    .build())
                    .then(Mono.fromRunnable {
                        log.debug("content-type: {}", exchange.request.headers.contentType)
                        val snapshot = exchange.attributes[ORIGINAL_REQUEST_BODY_SNAPSHOT] as? String
                        log.debug("peek request body snapshot:\n{}", snapshot)
                    })
        }

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

    inner class PeekedServerHttpRequest(private val exchange: ServerWebExchange, private val factoryConfig: Config) : ServerHttpRequestDecorator(exchange.request) {

        override fun getBody(): Flux<DataBuffer> {
            val flux = Flux.from(delegate.body)
                    .replay()
                    .autoConnect()

            flux.take(1).subscribe { buffer: DataBuffer ->
                val contentType: MediaType? = headers.contentType
                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_REQUEST_BODY_SNAPSHOT] = snapshot
                }
            }
            return flux
        }

    }

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

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