websocket

// WebSocket 配置类
@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(new ImportWebSocketHandler(), "/import-progress").setAllowedOrigins("*");
    }
}

// WebSocket 处理类
@Component
public class ImportWebSocketHandler extends TextWebSocketHandler {
    private Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        sessions.put(session.getId(), session);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        sessions.remove(session.getId());
    }

    public void sendProgress(String sessionId, String progress) throws IOException {
        WebSocketSession session = sessions.get(sessionId);
        if (session != null && session.isOpen()) {
            session.sendMessage(new TextMessage(progress));
        }
    }

    public void closeSession(String sessionId) throws IOException {
        WebSocketSession session = sessions.get(sessionId);
        if (session != null && session.isOpen()) {
            session.close();
        }
    }
}

@Path("/import")
@Produces(MediaType.APPLICATION_JSON)
@Component
public class ImportResource {
    @Autowired
    private ImportWebSocketHandler webSocketHandler;
    
    @Autowired
    private ExcelReadListener excelReadListener;

    @POST
    @Path("/upload")
    @Consumes(MediaType.MULTIPART_FORM_DATA)
    public Response uploadFile(@FormDataParam("file") InputStream fileInputStream,
                             @FormDataParam("file") FormDataContentDisposition fileDetail,
                             @HeaderParam("session-id") String sessionId) {
        try {
            // 检查文件类型
            String fileName = fileDetail.getFileName().toLowerCase();
            if (!fileName.endsWith(".xls") && !fileName.endsWith(".xlsx")) {
                return Response.status(Response.Status.BAD_REQUEST)
                    .entity("{\"status\": \"error\", \"message\": \"Please upload an Excel file (.xls or .xlsx)\"}")
                    .build();
            }

            // 获取总行数（需要先读取一次文件）
            long totalRows = getExcelRowCount(fileInputStream);
            fileInputStream.reset(); // 重置输入流
            
            // 设置监听器参数
            excelReadListener.setSessionId(sessionId);
            excelReadListener.setTotalRows(totalRows);

            // 使用 EasyExcel 读取文件
            EasyExcel.read(fileInputStream, ExcelData.class, excelReadListener)
                .sheet()
                .doRead();

            return Response.ok("{\"status\": \"success\"}").build();
        } catch (Exception e) {
            return Response.serverError()
                .entity("{\"status\": \"error\", \"message\": \"" + e.getMessage() + "\"}")
                .build();
        }
    }

    private long getExcelRowCount(InputStream inputStream) {
        AtomicLong rowCount = new AtomicLong(0);
        EasyExcel.read(inputStream, new NoModelDataListener() {
            @Override
            public void invoke(Object data, AnalysisContext context) {
                rowCount.incrementAndGet();
            }
        }).sheet().doRead();
        return rowCount.get();
    }
}

@Component
public class ExcelReadListener implements ReadListener<ExcelData> {
    @Autowired
    private ImportWebSocketHandler webSocketHandler;
    
    private String sessionId;
    private long totalRows;
    private long currentRow = 0;
    
    public void setSessionId(String sessionId) {
        this.sessionId = sessionId;
    }
    
    public void setTotalRows(long totalRows) {
        this.totalRows = totalRows;
    }

    @Override
    public void invoke(ExcelData data, AnalysisContext context) {
        currentRow++;
        try {
            // 处理每行数据
            processRow(data);
            
            // 计算并发送进度
            double progress = (double) currentRow / totalRows * 100;
            webSocketHandler.sendProgress(sessionId, String.format("%.2f%%", progress));
        } catch (IOException e) {
            throw new RuntimeException("Failed to send progress: " + e.getMessage());
        }
    }

    @Override
    public void doAfterAllAnalysed(AnalysisContext context) {
        try {
            webSocketHandler.sendProgress(sessionId, "100%");
            webSocketHandler.closeSession(sessionId);
        } catch (IOException e) {
            throw new RuntimeException("Failed to send final progress: " + e.getMessage());
        }
    }

    private void processRow(ExcelData data) {
        // 在这里实现具体的行处理逻辑
        System.out.println("Processing: " + data);
    }
}