package com.eactive.eai.adapter.http.dynamic.wrapper;

import io.undertow.server.HandlerWrapper;
import io.undertow.server.HttpHandler;
import io.undertow.server.HttpServerExchange;
import io.undertow.server.ServerConnection;
import io.undertow.util.AttachmentKey;
import io.undertow.util.Headers;

public class KeepAliveTrueWrapper implements HandlerWrapper {
	private final AttachmentKey<Integer> MAX_COUNT = AttachmentKey.create(Integer.class);
	private int keepAliveTimeout = 15;
	private int maxRequestCount = -1; // 최대 요청 횟수 설정

	public int getKeepAliveTimeout() {
		return keepAliveTimeout;
	}

	public void setKeepAliveTimeout(int keepAliveTimeout) {
		this.keepAliveTimeout = keepAliveTimeout;
	}

	public int getMaxRequestCount() {
		return maxRequestCount;
	}

	public void setMaxRequestCount(int maxRequestCount) {
		this.maxRequestCount = maxRequestCount;
	}

	@Override
	public HttpHandler wrap(HttpHandler handler) {
		return new HttpHandler() {
			@Override
			public void handleRequest(HttpServerExchange exchange) throws Exception {
				// max requst count 가 제한이 없을경우
				if (maxRequestCount == -1) {
					handler.handleRequest(exchange);
					exchange.getResponseHeaders().put(Headers.KEEP_ALIVE, "timeout=" + keepAliveTimeout);
					return;
				}

				// 현재 연결의 요청 횟수 가져오기
				ServerConnection con = exchange.getConnection();
				int count = 0;
				try {
					count = con.getAttachment(MAX_COUNT);
				} catch (Exception e) {

				}
				count++;

				con.putAttachment(MAX_COUNT, count);
				// 최대 요청 횟수를 초과한 경우 연결을 종료하고 로그를 출력
				if (count > maxRequestCount) {
					exchange.getConnection().close();
					// System.out.println("Connection closed due to exceeding max requests limit.");
					return;
				}
				handler.handleRequest(exchange);
				if (count > maxRequestCount - 1) {
					exchange.getResponseHeaders().put(Headers.CONNECTION, "close");
				}
				exchange.getResponseHeaders().put(Headers.KEEP_ALIVE,
						"timeout=" + keepAliveTimeout + ", max=" + (maxRequestCount - count));
			}
		};
	}
}
