package org.dromara.ai.websocket;

import lombok.RequiredArgsConstructor;
import org.dromara.ai.domain.dto.AiChatRequest;
import org.dromara.ai.domain.dto.AiChatResponse;
import org.dromara.ai.service.AiChatService;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.nio.charset.StandardCharsets;

@Component
@RequiredArgsConstructor
public class ChatWebSocketHandler implements WebSocketHandler {

    private final AiChatService aiChatService;

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) {
        try {
            String payload;
            if (message instanceof TextMessage) {
                payload = ((TextMessage) message).getPayload();
            } else if (message instanceof BinaryMessage) {
                payload = new String(((BinaryMessage) message).getPayload().array(), StandardCharsets.UTF_8);
            } else {
                return;
            }
            AiChatRequest req = org.dromara.common.json.utils.JsonUtils.parseObject(payload, AiChatRequest.class);
            aiChatService.chatStream(req, new AiChatService.StreamCallback() {
                @Override
                public void onChunk(String chunk) {
                    try {
                        session.sendMessage(new TextMessage(chunk));
                    } catch (Exception ignored) {
                    }
                }

                @Override
                public void onComplete(AiChatResponse response) {
                    try {
                        session.sendMessage(new TextMessage("[DONE]"));
                    } catch (Exception ignored) {
                    }
                }

                @Override
                public void onError(Throwable error) {
                    try {
                        session.sendMessage(new TextMessage("[ERROR] " + error.getMessage()));
                        session.close(CloseStatus.SERVER_ERROR);
                    } catch (Exception ignored) {
                    }
                }
            });
        } catch (Exception e) {
            try {
                session.sendMessage(new TextMessage("[ERROR] " + e.getMessage()));
                session.close(CloseStatus.SERVER_ERROR);
            } catch (Exception ignored) {
            }
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) {
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) {
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }
}
