package com.shzg.project.worn.websocket.endpoint;

import com.shzg.project.system.domain.SysDept;
import com.shzg.project.system.mapper.SysDeptMapper;
import com.shzg.project.worn.websocket.domain.WsUserInfo;
import com.shzg.project.worn.websocket.manager.WebSocketSessionManager;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

@Slf4j
@Component
@ServerEndpoint("/ws")
public class WornWebSocketServer {

    private static WebSocketSessionManager sessionManager;
    private static SysDeptMapper deptMapper;

    @Autowired
    public void setSessionManager(WebSocketSessionManager manager) {
        WornWebSocketServer.sessionManager = manager;
    }

    @Autowired
    public void setDeptMapper(SysDeptMapper mapper) {
        WornWebSocketServer.deptMapper = mapper;
    }

    @OnOpen
    public void onOpen(Session session) {

        log.info("[WebSocket] 连接建立 sessionId={}", session.getId());

        String query = session.getQueryString();

        if (query == null || !query.contains("deptId")) {
            closeSession(session, "missing deptId");
            return;
        }

        Long deptId;

        try {
            deptId = parseDeptId(query);
        } catch (Exception e) {
            closeSession(session, "invalid deptId");
            return;
        }

        List<SysDept> deptList = deptMapper.selectDeptAndChildren(deptId, null);

        Set<Long> deptIds = new HashSet<>();
        for (SysDept dept : deptList) {
            deptIds.add(dept.getDeptId());
        }

        // ===== 构造用户（当前写死）=====
        WsUserInfo userInfo = new WsUserInfo();
        userInfo.setUserId(1L);
        userInfo.setUserName("test");
        userInfo.setAdmin(false);
        userInfo.setDeptIds(deptIds);

        // ===== 获取IP（新增）=====
        String ip = getIp(session);

        // 存入 session（给 manager 用）
        session.getUserProperties().put("ip", ip);

        // ===== 注册连接（关键：接入限流）=====
        boolean success = sessionManager.register(session, userInfo);

        if (!success) {
            log.warn("[WebSocket] 连接被拒绝（限流）sessionId={}, ip={}", session.getId(), ip);
            closeSession(session, "too many connections");
            return;
        }

        log.info("[WebSocket] 注册成功 sessionId={}, ip={}, deptIds={}",
                session.getId(), ip, deptIds);
    }

    @OnClose
    public void onClose(Session session) {
        sessionManager.remove(session);
    }

    @OnError
    public void onError(Session session, Throwable error) {
        if (session != null) {
            sessionManager.remove(session);
        }
    }

    private Long parseDeptId(String query) {

        String[] params = query.split("&");

        for (String param : params) {
            if (param.startsWith("deptId=")) {
                return Long.parseLong(param.substring(7));
            }
        }

        throw new IllegalArgumentException("deptId not found");
    }

    /**
     * 获取客户端IP
     */
    private String getIp(Session session) {

        try {
            // 标准方式（部分容器支持）
            Object addr = session.getUserProperties().get("javax.websocket.endpoint.remoteAddress");

            if (addr instanceof InetSocketAddress) {
                return ((InetSocketAddress) addr).getAddress().getHostAddress();
            }

        } catch (Exception ignored) {}

        // fallback
        return "unknown";
    }

    private void closeSession(Session session, String reason) {
        try {
            session.close(new CloseReason(
                    CloseReason.CloseCodes.CANNOT_ACCEPT,
                    reason
            ));
        } catch (IOException e) {
            log.error("关闭失败", e);
        }
    }
}