package top.baogutang.strategy.infra;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import top.baogutang.strategy.commons.enums.Timeframe;
import top.baogutang.strategy.models.Kline;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;

@Slf4j
public class OkxWebSocketClient {

    private final OkHttpClient http;

    private final ObjectMapper mapper;

    private WebSocket ws;

    private final String URL = "wss://ws.okx.com:8443/ws/v5/business";

    public OkxWebSocketClient(OkHttpClient http, ObjectMapper mapper) {
        this.http = http;
        this.mapper = mapper;
    }

    public static class WsKlineEvent {
        private final String instId;
        private final Timeframe timeframe;
        private final Kline kline;

        public WsKlineEvent(String instId, Timeframe timeframe, Kline kline) {
            this.instId = instId;
            this.timeframe = timeframe;
            this.kline = kline;
        }

        public String getInstId() {
            return instId;
        }

        public Timeframe getTimeframe() {
            return timeframe;
        }

        public Kline getKline() {
            return kline;
        }
    }

    public void connect(final Map<String, List<Timeframe>> symbolMap,
                        final Consumer<WsKlineEvent> consumer) {
        Request req = new Request.Builder().url(URL).build();
        ws = http.newWebSocket(req, new WebSocketListener() {
            @Override
            public void onOpen(WebSocket webSocket, Response response) {
                log.info("[WS] onOpen code={} message={}", response.code(), response.message());
                List<Map<String, String>> args = new ArrayList<>();
                for (Map.Entry<String, List<Timeframe>> e : symbolMap.entrySet()) {
                    String inst = e.getKey();
                    for (Timeframe tf : e.getValue()) {
                        Map<String, String> arg = new HashMap<>();
                        arg.put("channel", tf.okxChannel);
                        arg.put("instId", inst);
                        args.add(arg);
                        log.info("[WS] will-subscribe channel={} inst={}", tf.okxChannel, inst);
                    }
                }
                Map<String, Object> sub = new HashMap<>();
                sub.put("op", "subscribe");
                sub.put("args", args);
                try {
                    String json = mapper.writeValueAsString(sub);
                    log.info("[WS] send subscribe payload={}", json);
                    webSocket.send(json);
                } catch (Exception ex) {
                    log.error("[WS] subscribe error", ex);
                }
            }

            @Override
            public void onMessage(WebSocket webSocket, String text) {
                log.debug("[WS] raw text={}", text);

                try {
                    JsonNode root = mapper.readTree(text);
                    // 订阅确认等事件
                    if (root.has("event")) {
                        log.info("[WS] event ack raw={}", text);
                        return;
                    }
                    if (!root.has("arg") || !root.has("data")) {
                        log.debug("[WS] skip no arg/data");
                        return;
                    }
                    JsonNode arg = root.get("arg");
                    String channel = arg.get("channel").asText();
                    String instId = arg.get("instId").asText();

                    Timeframe tf = null;
                    for (Timeframe t : Timeframe.values()) {
                        if (t.okxChannel.equals(channel)) {
                            tf = t;
                            break;
                        }
                    }
                    if (tf == null) {
                        log.warn("[WS] unknown channel={} inst={}", channel, instId);
                        return;
                    }

                    for (JsonNode arr : root.get("data")) {
                        // OKX K线数组字段，请核对文档；这里你目前只解析前 6 个
                        long ts = arr.get(0).asLong(); // 开始时间
                        double o = arr.get(1).asDouble();
                        double h = arr.get(2).asDouble();
                        double l = arr.get(3).asDouble();
                        double c = arr.get(4).asDouble();
                        double v = arr.get(5).asDouble();
                        // 建议用 confirm 标志（最后一个字段），而不是本地时间推断
                        boolean finalBar = false;
                        if (arr.size() >= 10) { // 根据 OKX 实际返回字段数调整
                            // 一般 confirm 在最后一个
                            JsonNode confirmNode = arr.get(arr.size() - 1);
                            if (confirmNode.isTextual()) {
                                finalBar = "1".equals(confirmNode.asText());
                            } else if (confirmNode.isInt() || confirmNode.isBoolean()) {
                                finalBar = confirmNode.asInt() == 1 || confirmNode.asBoolean();
                            }
                        } else {
                            // 保留你的旧逻辑作为兜底
                            finalBar = System.currentTimeMillis() >= ts + tf.millis;
                        }

                        Kline k = new Kline(ts, o, h, l, c, v, ts + tf.millis - 1, finalBar);
//                        log.info("[WS] DATA inst={} tf={} o={} h={} l={} c={} v={} finalBar={} size={}",
//                                instId, tf, o, h, l, c, v, finalBar, arr.size());
                        consumer.accept(new WsKlineEvent(instId, tf, k));
                    }
                } catch (Exception e) {
                    log.error("[WS] onMessage parse error text={}", text, e);
                }
            }

            @Override
            public void onClosing(WebSocket webSocket, int code, String reason) {
                log.warn("[WS] onClosing code={} reason={}", code, reason);
            }

            @Override
            public void onClosed(WebSocket webSocket, int code, String reason) {
                log.warn("[WS] onClosed code={} reason={}", code, reason);
            }

            @Override
            public void onFailure(WebSocket webSocket, Throwable t, Response response) {
                log.error("[WS] onFailure respCode={} respMsg={}",
                        response == null ? null : response.code(),
                        response == null ? null : response.message(), t);
            }
        });

    }
}
