package com.grigowashere.aismap.ble.hub;

import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

/**
 * Reassembles chunked logical messages keyed by (session_msg_id, msg_type).
 *
 * Payload may be JSON or MessagePack depending on the protocol-version byte of
 * the incoming frames (see {@link AisHubConstants#PROTO_VERSION_JSON} /
 * {@link AisHubConstants#PROTO_VERSION_MSGPACK}). The assembler stores the
 * version of the first frame of an assembly and returns it to the caller when
 * the message is complete; the caller then chooses the right decoder.
 */
public class AisHubChunkAssembler {

    public static final class FeedStatus {
        /**
         * Non-null only when the whole message is complete.
         * Historically this was a JSON string; it's kept for backward
         * compatibility with callers that only ever saw JSON on the wire.
         * When the assembly's proto version is MessagePack, {@link #json} is
         * {@code null} even at completion — use {@link #payload} +
         * {@link #protocolVersion} instead.
         */
        public final String json;
        /**
         * Full reassembled payload bytes (always non-null when the message is
         * complete, regardless of encoding). Callers that support both JSON
         * and MessagePack should decode from this field and ignore
         * {@link #json}.
         */
        public final byte[] payload;
        /** Protocol-version byte of the reassembled message (matches {@link AisHubConstants}). */
        public final int protocolVersion;
        public final int received;
        public final int chunkCount;
        /** True if we detected mismatch and reset assembly state. */
        public final boolean wasReset;

        FeedStatus(String json, byte[] payload, int protocolVersion,
                   int received, int chunkCount, boolean wasReset) {
            this.json = json;
            this.payload = payload;
            this.protocolVersion = protocolVersion;
            this.received = received;
            this.chunkCount = chunkCount;
            this.wasReset = wasReset;
        }
    }

    private static final class Key {
        final int sessionId;
        final int msgType;

        Key(int sessionId, int msgType) {
            this.sessionId = sessionId;
            this.msgType = msgType;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (!(o instanceof Key)) return false;
            Key key = (Key) o;
            return sessionId == key.sessionId && msgType == key.msgType;
        }

        @Override
        public int hashCode() {
            return 31 * sessionId + msgType;
        }
    }

    private static final class Assembly {
        final int chunkCount;
        final byte[][] parts;
        int received;
        // Protocol version taken from the first frame of this assembly.
        // We pin it so that a rogue frame with a mismatched version cannot
        // silently change how we decode the combined payload.
        int protocolVersion = -1;

        Assembly(int chunkCount) {
            this.chunkCount = chunkCount;
            this.parts = new byte[chunkCount][];
        }
    }

    private final Map<Key, Assembly> pending = new HashMap<>();

    /**
     * Feed one frame; returns complete UTF-8 JSON string or null if more chunks needed.
     * <p>
     * <b>Deprecated for mixed encodings:</b> MessagePack payloads will return
     * {@code null} here even when the message is complete. Use
     * {@link #feedStatus(AisHubFrame)} and inspect
     * {@link FeedStatus#payload} + {@link FeedStatus#protocolVersion} instead.
     */
    public String feed(AisHubFrame frame) {
        FeedStatus st = feedStatus(frame);
        return st != null ? st.json : null;
    }

    /**
     * Same as {@link #feed(AisHubFrame)}, but also returns progress info for logging
     * and the raw reassembled payload + protocol version on completion.
     */
    public FeedStatus feedStatus(AisHubFrame frame) {
        if (frame.chunkCount <= 0 || frame.chunkIndex < 0 || frame.chunkIndex >= frame.chunkCount) {
            return new FeedStatus(null, null, frame.protocolVersion, 0, frame.chunkCount, false);
        }
        Key key = new Key(frame.sessionMsgId, frame.msgType);
        Assembly a = pending.get(key);
        boolean wasReset = false;
        if (a == null) {
            a = new Assembly(frame.chunkCount);
            a.protocolVersion = frame.protocolVersion;
            pending.put(key, a);
        } else if (a.chunkCount != frame.chunkCount || a.protocolVersion != frame.protocolVersion) {
            // Either the server changed its mind about the chunk count (unlikely
            // but defensively handled) or the encoding — treat as a brand new
            // assembly. Can happen if a previous message was lost and we're now
            // seeing the start of the next one under the same (sid, msg_type).
            pending.remove(key);
            a = new Assembly(frame.chunkCount);
            a.protocolVersion = frame.protocolVersion;
            pending.put(key, a);
            wasReset = true;
        }
        if (a.parts[frame.chunkIndex] == null) {
            a.received++;
        }
        a.parts[frame.chunkIndex] = Arrays.copyOf(frame.payload, frame.payload.length);
        if (a.received < a.chunkCount) {
            return new FeedStatus(null, null, a.protocolVersion, a.received, a.chunkCount, wasReset);
        }
        int total = 0;
        for (byte[] p : a.parts) {
            if (p != null) total += p.length;
        }
        byte[] out = new byte[total];
        int pos = 0;
        for (byte[] p : a.parts) {
            if (p != null) {
                System.arraycopy(p, 0, out, pos, p.length);
                pos += p.length;
            }
        }
        pending.remove(key);
        // Build the JSON string only for legacy JSON payloads, to keep old
        // callers (that read FeedStatus.json directly) working as-is.
        String jsonStr = (a.protocolVersion == AisHubConstants.PROTO_VERSION_JSON)
                ? new String(out, StandardCharsets.UTF_8)
                : null;
        return new FeedStatus(jsonStr, out, a.protocolVersion, a.received, a.chunkCount, wasReset);
    }

    public void clear() {
        pending.clear();
    }
}
