package gc.mda.signal_batch.global.websocket.interceptor;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.stereotype.Component;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;


@Slf4j
@Component
public class TrackQueryInterceptor implements ChannelInterceptor {

    // 세션별 활성 쿼리 수 추적
    private final ConcurrentHashMap<String, AtomicInteger> sessionQueries = new ConcurrentHashMap<>();

    @Value("${websocket.query.max-per-session:3}")
    private int maxQueriesPerSession;

    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
        
        if (accessor != null && StompCommand.SEND.equals(accessor.getCommand())) {
            String destination = accessor.getDestination();
            String sessionId = accessor.getSessionId();
            
            // 궤적 쿼리 요청 확인
            if ("/app/tracks/query".equals(destination)) {
                AtomicInteger queryCount = sessionQueries.computeIfAbsent(
                    sessionId, k -> new AtomicInteger(0)
                );
                
                if (queryCount.get() >= maxQueriesPerSession) {
                    log.warn("Session {} exceeded max queries limit ({})", sessionId, maxQueriesPerSession);
                    throw new IllegalStateException(
                        "Maximum concurrent queries per session exceeded: " + maxQueriesPerSession);
                }
                
                queryCount.incrementAndGet();
                log.info("Session {} started query. Active queries: {}", 
                        sessionId, queryCount.get());
            }
        }
        
        return message;
    }

    @Override
    public void afterSendCompletion(Message<?> message, MessageChannel channel, 
                                   boolean sent, Exception ex) {
        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(
            message, StompHeaderAccessor.class
        );
        
        if (accessor != null && StompCommand.DISCONNECT.equals(accessor.getCommand())) {
            String sessionId = accessor.getSessionId();
            sessionQueries.remove(sessionId);
            log.info("Session {} disconnected, cleared query count", sessionId);
        }
    }
    
    /**
     * 쿼리 완료/실패 시 세션 카운트 감소 + 글로벌 슬롯 반환
     * 스트리밍 서비스의 finally 블록에서 호출
     */
    public void releaseQuery(String sessionId) {
        AtomicInteger count = sessionQueries.get(sessionId);
        if (count != null && count.get() > 0) {
            int remaining = count.decrementAndGet();
            log.debug("Session {} query released. Remaining: {}", sessionId, remaining);
        }
    }
}