import com.fasterxml.jackson.annotation.JsonInclude

@JsonInclude(JsonInclude.Include.NON_NULL)
class ApiResponse<T> (
    val data: T?,
    var result: Boolean = true,
    var message: String? = null
) {

    companion object {
        fun success(data: Any): ApiResponse<Any> {
            return ApiResponse(data)
        }

        fun failed(message: String? = null): ApiResponse<Any> {
            return ApiResponse(null, false, message)
        }
    }
}

@RestController
@RequestMapping("/boards")
class BoardController(
    val boardService: BoardService
) {

    @PostMapping
    fun create(@RequestBody boardRequestDto: BoardRequestDto) =
        boardService.create(boardRequestDto).map { ApiResponse(it) }

    @PutMapping("/{boardId}")
    fun update(
        @RequestBody boardUpdateDto: BoardUpdateDto,
        @PathVariable("boardId") boardId: Long
    ) =
        boardService.update(boardId, boardUpdateDto).map { ApiResponse(it) }

    @GetMapping("/{boardId}")
    fun getOne(@PathVariable("boardId") boardId: Long) = boardService.get(boardId).map { ApiResponse(it) }

    @GetMapping
    fun getAll() = boardService.getAll().map { ApiResponse(it) }
}

import me.sup2is.kotlinreactiveboard.api.controller.model.ApiResponse
import org.springframework.http.codec.HttpMessageWriter
import org.springframework.web.reactive.HandlerResult
import org.springframework.web.reactive.accept.RequestedContentTypeResolver
import org.springframework.web.reactive.result.method.annotation.ResponseBodyResultHandler
import org.springframework.web.server.ServerWebExchange
import reactor.core.publisher.Flux
import reactor.core.publisher.Mono
import reactor.kotlin.core.publisher.toMono

class ResponseWrapper(
    writers: MutableList<HttpMessageWriter<*>>,
    resolver: RequestedContentTypeResolver
) : ResponseBodyResultHandler(writers, resolver) {

    override fun supports(result: HandlerResult): Boolean {
        val isSupportTypes = result.returnType.resolve() == Mono::class.java ||
            result.returnType.resolve() == Flux::class.java
        val isAlreadyResponse = result.returnType.generics[0] == ApiResponse::class
        return isSupportTypes && !isAlreadyResponse
    }

    override fun handleResult(exchange: ServerWebExchange, result: HandlerResult): Mono<Void> {

        val body = when (val value = result.returnValue) {
            is Mono<*> -> value
            is Flux<*> -> value.collectList()
            else -> throw ClassCastException("The \"body\" should be Mono<*> or Flux<*>!")
        }.map {
            ApiResponse.success(it)
        }.onErrorResume {
            ApiResponse.failed(it.message).toMono()
        }

        val returnTypeSource = result.returnTypeSource

        return writeBody(body, returnTypeSource, exchange)
    }
}

@Configuration
class WebConfig {

    @Bean
    fun requestWrapper(
        serverCodecConfigurer: ServerCodecConfigurer,
        requestedContentTypeResolver: RequestedContentTypeResolver
    ): ResponseWrapper {
        return ResponseWrapper(serverCodecConfigurer.writers, requestedContentTypeResolver)
    }
}

@WebFluxTest(BoardController::class)
@Import(WebConfig::class)
class ResponseWrapperTest {

    @Autowired
    lateinit var webTestClient: WebTestClient

    @MockBean
    lateinit var boardService: BoardService

    @Test
    fun `API가 성공하면 ApiResponse에 wrapping한다`() {

        // given
        given(boardService.get(anyLong()))
            .willReturn(Board().toMono())

        // when & then
        webTestClient.get()
            .uri("/boards/{boardId}", 1)
            .exchange()
            .expectStatus().isOk
            .expectBody(ApiResponse::class.java)
            .consumeWith {
                assertThat(it.responseBody!!.result).isTrue
                assertThat(it.responseBody!!.data).isNotNull
            }
    }

    @Test
    fun `API가 성공하면 ApiResponse에 wrapping한다2`() {

        // given
        given(boardService.getAll())
            .willReturn(Flux.just(Board(), Board()))

        // when & then
        webTestClient.get()
            .uri("/boards")
            .exchange()
            .expectStatus().isOk
            .expectBody(ApiResponse::class.java)
            .consumeWith {
                assertThat(it.responseBody!!.result).isTrue
                assertThat(it.responseBody!!.data).isNotNull
            }
    }
}

// DispatcherHandler.java

	@Override
	public Mono<Void> handle(ServerWebExchange exchange) {
		if (this.handlerMappings == null) {
			return createNotFoundError();
		}
		if (CorsUtils.isPreFlightRequest(exchange.getRequest())) {
			return handlePreFlight(exchange);
		}
		return Flux.fromIterable(this.handlerMappings)
				.concatMap(mapping -> mapping.getHandler(exchange))
				.next()
				.switchIfEmpty(createNotFoundError())
				.flatMap(handler -> invokeHandler(exchange, handler))
				.flatMap(result -> handleResult(exchange, result));
	}

	private Mono<Void> handleResult(ServerWebExchange exchange, HandlerResult result) {
		return getResultHandler(result).handleResult(exchange, result)
				.checkpoint("Handler " + result.getHandler() + " [DispatcherHandler]")
				.onErrorResume(ex ->
						result.applyExceptionHandler(ex).flatMap(exResult -> {
							String text = "Exception handler " + exResult.getHandler() +
									", error=\"" + ex.getMessage() + "\" [DispatcherHandler]";
							return getResultHandler(exResult).handleResult(exchange, exResult).checkpoint(text);
						}));
	}

	private HandlerResultHandler getResultHandler(HandlerResult handlerResult) {
		if (this.resultHandlers != null) {
			for (HandlerResultHandler resultHandler : this.resultHandlers) {
				if (resultHandler.supports(handlerResult)) {
					return resultHandler;
				}
			}
		}
		throw new IllegalStateException("No HandlerResultHandler for " + handlerResult.getReturnValue());
	}