// Copyright (C) 2026 Whiterun LLC,
// This software is licensed under the GNU Lesser General Public License (LGPL), version 3.0 or later.
// A copy of the license can be found in the LICENSE file or at https://www.gnu.org/licenses/lgpl-3.0.html

import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
import { ChunkAssembler, CHUNK_TIMEOUT_MS } from "./chunk-assembler.js";
import { type ChunkRelayMessage, RelayMsgAction } from "./protocols/hdwalletv1.js";

function makeChunk(
  chunk_id: string,
  chunk_index: number,
  chunk_total: number,
  chunk_data: string,
): ChunkRelayMessage {
  return {
    action: RelayMsgAction.ChunkRelay,
    chunk_id,
    chunk_index,
    chunk_total,
    chunk_data,
    time: Math.floor(Date.now() / 1000),
  };
}

function splitMessage(message: object, chunkSize: number, chunk_id: string): ChunkRelayMessage[] {
  const json = JSON.stringify(message);
  const chunks: ChunkRelayMessage[] = [];
  const total = Math.ceil(json.length / chunkSize);
  for (let i = 0; i < total; i++) {
    chunks.push(makeChunk(chunk_id, i, total, json.slice(i * chunkSize, (i + 1) * chunkSize)));
  }
  return chunks;
}

describe("ChunkAssembler", () => {
  let assembler: ChunkAssembler;

  beforeEach(() => {
    vi.useFakeTimers();
    assembler = new ChunkAssembler();
  });

  afterEach(() => {
    vi.useRealTimers();
  });

  describe("single-fragment messages", () => {
    it("reassembles a single-fragment message immediately", () => {
      const original = { action: RelayMsgAction.SignCancel, sequence: 1, time: 1000 };
      const [chunk] = splitMessage(original, 10000, "id1");

      const result = assembler.addChunk(chunk);

      expect(result).toEqual(original);
    });

    it("clears the pending buffer after assembly", () => {
      const original = { action: RelayMsgAction.SignCancel, sequence: 1, time: 1000 };
      const [chunk] = splitMessage(original, 10000, "id1");

      assembler.addChunk(chunk);

      expect(assembler.pendingCount()).toBe(0);
    });
  });

  describe("multi-fragment messages", () => {
    it("returns null for each fragment until the last", () => {
      const original = { action: RelayMsgAction.SignCancel, sequence: 7, time: 1000 };
      const chunks = splitMessage(original, 5, "id2");
      expect(chunks.length).toBeGreaterThan(1);

      for (let i = 0; i < chunks.length - 1; i++) {
        expect(assembler.addChunk(chunks[i])).toBeNull();
      }
    });

    it("returns the reassembled message on the last fragment", () => {
      const original = { action: RelayMsgAction.SignCancel, sequence: 7, time: 1000 };
      const chunks = splitMessage(original, 5, "id3");

      let result = null;
      for (const chunk of chunks) {
        result = assembler.addChunk(chunk);
      }

      expect(result).toEqual(original);
    });

    it("reassembles correctly when fragments arrive out of order", () => {
      const original = { action: RelayMsgAction.SignCancel, sequence: 42, time: 2000 };
      const chunks = splitMessage(original, 3, "id4");
      expect(chunks.length).toBeGreaterThan(1);

      // Reverse order
      const reversed = [...chunks].reverse();
      let result = null;
      for (const chunk of reversed) {
        result = assembler.addChunk(chunk);
      }

      expect(result).toEqual(original);
    });

    it("handles three fragments out of order (shuffle)", () => {
      const original = {
        action: RelayMsgAction.DappReady,
        supported_protocols: ["hdwalletv1"],
        wallet_discovered: false,
        time: 3000,
      };
      const chunks = splitMessage(original, 10, "id5");
      expect(chunks.length).toBeGreaterThanOrEqual(3);

      // Send middle first, then last, then first
      const [c0, c1, c2, ...rest] = chunks;
      const reordered = [c1, c2, c0, ...rest];
      let result = null;
      for (const chunk of reordered) {
        result = assembler.addChunk(chunk);
      }

      expect(result).toEqual(original);
    });

    it("reassembles a large payload split into many fragments", () => {
      const original = {
        action: RelayMsgAction.SignTransactionRequest,
        sequence: 999,
        time: 5000,
        // Simulate a large payload
        transaction: { data: "x".repeat(5000) },
        inputPaths: [],
      };
      const chunks = splitMessage(original, 500, "id6");
      expect(chunks.length).toBeGreaterThan(5);

      let result = null;
      for (const chunk of chunks) {
        result = assembler.addChunk(chunk);
      }

      expect(result).toEqual(original);
    });
  });

  describe("independent chunk_ids", () => {
    it("tracks multiple concurrent assemblies independently", () => {
      const msgA = { action: RelayMsgAction.SignCancel, sequence: 1, time: 1000 };
      const msgB = { action: RelayMsgAction.SignCancel, sequence: 2, time: 2000 };
      const chunksA = splitMessage(msgA, 5, "idA");
      const chunksB = splitMessage(msgB, 5, "idB");

      // Interleave: A0, B0, A1, B1, ...
      let resultA = null;
      let resultB = null;
      const maxLen = Math.max(chunksA.length, chunksB.length);
      for (let i = 0; i < maxLen; i++) {
        if (i < chunksA.length) resultA = assembler.addChunk(chunksA[i]);
        if (i < chunksB.length) resultB = assembler.addChunk(chunksB[i]);
      }

      expect(resultA).toEqual(msgA);
      expect(resultB).toEqual(msgB);
    });

    it("completing one assembly does not affect pending assemblies", () => {
      const msgA = { action: RelayMsgAction.SignCancel, sequence: 1, time: 1000 };
      const msgB = { action: RelayMsgAction.SignCancel, sequence: 2, time: 2000 };
      const [chunkA] = splitMessage(msgA, 10000, "idA");
      const chunksB = splitMessage(msgB, 5, "idB");

      // Complete A immediately
      assembler.addChunk(chunkA);

      // B is still in progress
      for (let i = 0; i < chunksB.length - 1; i++) {
        assembler.addChunk(chunksB[i]);
      }

      expect(assembler.pendingCount()).toBe(1);

      // Complete B
      const result = assembler.addChunk(chunksB[chunksB.length - 1]);
      expect(result).toEqual(msgB);
      expect(assembler.pendingCount()).toBe(0);
    });
  });

  describe("error cases", () => {
    it("returns null and discards buffer on inconsistent chunk_total", () => {
      const chunk1 = makeChunk("id7", 0, 3, "part1");
      const chunk2 = makeChunk("id7", 1, 4, "part2"); // wrong total

      assembler.addChunk(chunk1);
      const result = assembler.addChunk(chunk2);

      expect(result).toBeNull();
      expect(assembler.pendingCount()).toBe(0);
    });

    it("returns null for invalid JSON", () => {
      const chunk = makeChunk("id8", 0, 1, "{not valid json}}}");
      const result = assembler.addChunk(chunk);

      expect(result).toBeNull();
    });

    it("returns null if assembled payload is not a ProtocolMessage", () => {
      const notAMessage = { foo: "bar" };
      const [chunk] = splitMessage(notAMessage, 10000, "id9");
      const result = assembler.addChunk(chunk);

      expect(result).toBeNull();
    });

    it("returns null for chunk_total of 0", () => {
      const chunk = makeChunk("id10", 0, 0, "data");
      const result = assembler.addChunk(chunk);

      expect(result).toBeNull();
    });
  });

  describe("timeout cleanup", () => {
    it("discards incomplete assembly after timeout", () => {
      const msgA = { action: RelayMsgAction.SignCancel, sequence: 1, time: 1000 };
      const chunks = splitMessage(msgA, 5, "idTimeout");
      expect(chunks.length).toBeGreaterThan(1);

      // Add all but the last fragment
      for (let i = 0; i < chunks.length - 1; i++) {
        assembler.addChunk(chunks[i]);
      }
      expect(assembler.pendingCount()).toBe(1);

      vi.advanceTimersByTime(CHUNK_TIMEOUT_MS + 1);

      expect(assembler.pendingCount()).toBe(0);
    });

    it("does not timeout before CHUNK_TIMEOUT_MS", () => {
      const msgA = { action: RelayMsgAction.SignCancel, sequence: 1, time: 1000 };
      const chunks = splitMessage(msgA, 5, "idTimeout2");

      for (let i = 0; i < chunks.length - 1; i++) {
        assembler.addChunk(chunks[i]);
      }

      vi.advanceTimersByTime(CHUNK_TIMEOUT_MS - 1);

      expect(assembler.pendingCount()).toBe(1);
    });

    it("clears timeout when assembly completes before timeout", () => {
      const msgA = { action: RelayMsgAction.SignCancel, sequence: 1, time: 1000 };
      const chunks = splitMessage(msgA, 5, "idTimeout3");

      for (const chunk of chunks) {
        assembler.addChunk(chunk);
      }
      expect(assembler.pendingCount()).toBe(0);

      // Advancing past timeout should not throw
      expect(() => vi.advanceTimersByTime(CHUNK_TIMEOUT_MS + 1)).not.toThrow();
    });

    it("after timeout, late arriving fragments start a fresh assembly", () => {
      const original = { action: RelayMsgAction.SignCancel, sequence: 1, time: 1000 };
      const chunks = splitMessage(original, 5, "idReuse");

      // Send first fragment, let it time out
      assembler.addChunk(chunks[0]);
      vi.advanceTimersByTime(CHUNK_TIMEOUT_MS + 1);
      expect(assembler.pendingCount()).toBe(0);

      // Re-send all fragments with same chunk_id — should reassemble fresh
      let result = null;
      for (const chunk of chunks) {
        result = assembler.addChunk(chunk);
      }
      expect(result).toEqual(original);
    });
  });

  describe("duplicate fragment index", () => {
    it("last write wins for duplicate index — still assembles correctly if content is the same", () => {
      const original = { action: RelayMsgAction.SignCancel, sequence: 1, time: 1000 };
      const chunks = splitMessage(original, 5, "idDup");
      expect(chunks.length).toBeGreaterThan(1);

      // Send first fragment twice (relay may deliver duplicates)
      assembler.addChunk(chunks[0]);
      assembler.addChunk(chunks[0]);

      let result = null;
      for (let i = 1; i < chunks.length; i++) {
        result = assembler.addChunk(chunks[i]);
      }

      // Assembly happens when fragment count equals chunk_total.
      // With duplicate index 0, we have total fragments but index 0 is stored once.
      // The assembly should still complete when index N-1 is added.
      if (result !== null) {
        expect(result).toEqual(original);
      }
      // (If size counting means assembly triggers early on the duplicate, result may be null here
      //  and the final fragment triggers it — either way is acceptable.)
    });
  });
});