package com.vetti.socket.agents.ai;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;

import jakarta.websocket.*;
import java.net.URI;
import java.util.Collections;
import java.util.List;
import java.util.Map;

public class ElevenLabsAgentEndpoint extends Endpoint {

    private Session session;
    private final ObjectMapper mapper = new ObjectMapper();

    @Override
    public void onOpen(Session session, EndpointConfig config) {
        this.session = session;
        System.out.println("WebSocket opened");

        session.addMessageHandler(String.class, this::onMessage);
    }

    @Override
    public void onClose(Session session, CloseReason closeReason) {
        System.out.println("WebSocket closed: " + closeReason);
    }

    @Override
    public void onError(Session session, Throwable thr) {
        thr.printStackTrace();
    }

    private void onMessage(String message) {
        try {
            JsonNode json = mapper.readTree(message);
            System.out.println("Received: " + json.toPrettyString());

            // 判断 turn 结束
            if ("agent_chat_response_part".equals(json.path("type").asText())) {
                JsonNode part = json.path("text_response_part");
                if ("stop".equals(part.path("type").asText())) {
                    System.out.println("Turn ended");
                }
            }

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public void sendTextMessage(String text) {
        if (session != null && session.isOpen()) {
            String msgJson = String.format("{\"type\":\"input_text\",\"text\":\"%s\"}", text);
            session.getAsyncRemote().sendText(msgJson);
        } else {
            throw new IllegalStateException("Session not open");
        }
    }

    public void connect(String agentId, String apiKey) throws Exception {
        WebSocketContainer container = ContainerProvider.getWebSocketContainer();

        // 配置 HTTP header
        ClientEndpointConfig.Configurator configurator = new ClientEndpointConfig.Configurator() {
            @Override
            public void beforeRequest(Map<String, List<String>> headers) {
                headers.put("xi-api-key", Collections.singletonList(apiKey));
            }
        };
        ClientEndpointConfig config = ClientEndpointConfig.Builder.create()
                .configurator(configurator)
                .build();

        String url = "wss://api.elevenlabs.io/v1/convai/conversation?agent_id=" + agentId;

        // 注意，这里是关键：返回 Session
        Session wsSession = container.connectToServer(this, config, URI.create(url));
        this.session = wsSession;
    }

    public void close() throws Exception {
        if (session != null) {
            session.close();
        }
    }
}

