package com.cmvr.framework.websocket.registry;

import com.cmvr.framework.websocket.annotation.WSMessageType;
import com.cmvr.framework.websocket.enums.WSMessageTypeEnum;
import com.cmvr.framework.websocket.handler.WSMessageHandler;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.Getter;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketSession;

import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.Map;

@Component
public class MessageHandlerRegistry implements ApplicationContextAware {
    private final Map<WSMessageTypeEnum, HandlerEntry<?, ?>> handlerMap = new HashMap<>();

    @Override
    public void setApplicationContext(ApplicationContext context) {
        Map<String, Object> beans = context.getBeansWithAnnotation(WSMessageType.class);
        for (Object bean : beans.values()) {
            if (bean instanceof WSMessageHandler) {
                WSMessageHandler<?> handler = (WSMessageHandler<?>) bean;
                WSMessageType anno = bean.getClass().getAnnotation(WSMessageType.class);
                handlerMap.put(anno.value(), new HandlerEntry<>(handler, anno.messageClass()));
            }
        }
    }

    public void handle(WebSocketSession session, String type, JsonNode payload) throws Exception {
        WSMessageTypeEnum typeEnum = WSMessageTypeEnum.fromValue(type);
        HandlerEntry<?, ?> entry = handlerMap.get(typeEnum);
        if (entry == null) {
            throw new IllegalArgumentException("未知消息类型: " + type);
        }

        Object message = new ObjectMapper().treeToValue(payload, entry.getMessageClass());
        entry.getHandleMethod().invoke(entry.getHandler(), session, message);
    }

    @Getter
    static class HandlerEntry<H extends WSMessageHandler<?>, T> {
        private final H handler;
        private final Class<T> messageClass;
        private final Method handleMethod;

        public HandlerEntry(H handler, Class<T> messageClass) {
            this.handler = handler;
            this.messageClass = messageClass;
            try {
                this.handleMethod = handler.getClass().getMethod("handle", WebSocketSession.class, messageClass);
            } catch (NoSuchMethodException e) {
                throw new IllegalStateException("找不到 handle 方法", e);
            }
        }

    }
}
