import { Type } from "@sinclair/typebox";
import { describe, expect, it } from "bun:test";
import { ExtendedPatch } from "../extended";
import { MultiplayerStateManager } from "../multiplayer";
import { ServerScripts } from "../multiplayerPolicies";
import { ServerScriptManager } from "../serverScripts";
import { createUnsafeQuickJSModule } from "../serverSimulator";
import { HistoryEntry } from "../stateManager";

const testEnvironment = {
  target: "test" as const,
  mode: "simulation" as const,
  label: "server-sim",
};

type TestState = { count: number };

describe("server prediction", () => {
  it("tracks and clears predictions when matching server invocation arrives", () => {
    const manager = new MultiplayerStateManager<TestState>({ count: 0 }, {});
    manager.connect();

    // Inject a predicted server step
    manager.applyServerPrediction({
      nextState: { count: 1 },
      invocation: 1,
      timestamp: 1100,
    });

    expect(manager.pendingCommits.length).toBe(1);
    const latestPrediction = manager.getLatestPrediction?.();
    expect(latestPrediction?.state.count).toBe(1);
    expect(latestPrediction?.serverInvocation?.invocation).toBe(1);

    const patches: ExtendedPatch[] = [
      { op: "replace", path: ["count"], value: 1 } as ExtendedPatch,
    ];

    manager.incomingMessages.push({
      type: "acceptPatch",
      id: "server-1",
      patches,
      serverInvocation: { invocation: 1, timestamp: 1100 },
    });

    manager.processIncomingMessages();

    // Prediction should be cleared and state committed
    expect(manager.pendingCommits.length).toBe(0);
    expect(manager.getConfirmedState().count).toBe(1);
  });

  it("updates clock skew from server invocation metadata", () => {
    const manager = new MultiplayerStateManager<TestState>(
      { count: 0 },
      { getNow: () => 1000 }
    );
    manager.connect();

    manager.incomingMessages.push({
      type: "acceptPatch",
      id: "skew-check",
      patches: [],
      serverInvocation: { invocation: 2, timestamp: 1500 },
    });

    manager.processIncomingMessages();

    const diagnostics = manager.getServerSimulationDiagnostics();
    expect(diagnostics.clockSkewMs).toBe(500);
    expect(
      diagnostics.invocations.find(
        (i) => (i.scriptId ?? "default") === "default"
      )?.invocation
    ).toBe(2);
  });

  it("keeps predictions for other scripts when only one script resolves", () => {
    const manager = new MultiplayerStateManager<TestState>({ count: 0 }, {});
    manager.connect();

    manager.applyServerPrediction({
      scriptId: "a",
      nextState: { count: 1 },
      invocation: 1,
      timestamp: 1000,
    });
    manager.applyServerPrediction({
      scriptId: "b",
      nextState: { count: 2 },
      invocation: 1,
      timestamp: 1000,
    });

    expect(manager.pendingCommits.length).toBe(2);

    manager.incomingMessages.push({
      type: "acceptPatch",
      id: "server-a",
      patches: [],
      serverInvocation: { scriptId: "a", invocation: 1, timestamp: 1000 },
    });
    manager.processIncomingMessages();

    expect(
      manager.pendingCommits.find(
        (c) => c.historyEntry.metadata.serverInvocation?.scriptId === "b"
      )
    ).toBeTruthy();
  });

  it("ignores predictions when disconnected", () => {
    const manager = new MultiplayerStateManager<TestState>({ count: 0 }, {});
    // not calling connect()

    manager.applyServerPrediction({
      scriptId: "a",
      nextState: { count: 1 },
      invocation: 1,
      timestamp: 1000,
    });

    expect(manager.pendingCommits.length).toBe(0);
  });

  it("clears prediction maps on reinitialize", () => {
    const manager = new MultiplayerStateManager<TestState>({ count: 0 }, {});
    manager.connect();

    manager.applyServerPrediction({
      nextState: { count: 1 },
      invocation: 1,
      timestamp: 1000,
    });

    expect(manager.pendingCommits.length).toBe(1);
    manager.reinitialize({ count: 0 });
    expect(manager.pendingCommits.length).toBe(0);
    expect(
      manager.getServerSimulationDiagnostics().pendingPredictions.length
    ).toBe(0);
  });

  it("attaches serverInvocation metadata to outgoing patch when present", () => {
    const manager = new MultiplayerStateManager<TestState>({ count: 0 }, {});
    manager.connect();

    const historyEntry: Required<
      HistoryEntry<TestState, typeof manager._metadataType>
    > = {
      undo: (s: TestState) => s,
      redo: (s: TestState) => s,
      metadata: {
        id: "m1",
        timestamp: Date.now(),
        serverInvocation: {
          scriptId: "s1",
          invocation: 1,
          timestamp: 1000,
          source: "client-sim",
        },
      },
      redoPatches: [],
      undoPatches: [],
    };

    manager._pushCommit({ count: 0 }, historyEntry, { addToHistory: false });

    const patchMsg = manager.outgoingMessages.find((m) => m.type === "patch");
    expect((patchMsg as any).serverInvocation?.scriptId).toBe("s1");
  });

  it("smooths clock skew over multiple updates", () => {
    let now = 1000;
    const manager = new MultiplayerStateManager<TestState>(
      { count: 0 },
      { getNow: () => now }
    );
    manager.connect();

    manager.incomingMessages.push({
      type: "acceptPatch",
      id: "skew1",
      patches: [],
      serverInvocation: { invocation: 1, timestamp: 1300 },
    });
    manager.processIncomingMessages();
    const skew1 = manager.getServerSimulationDiagnostics().clockSkewMs;

    now = 2000;
    manager.incomingMessages.push({
      type: "acceptPatch",
      id: "skew2",
      patches: [],
      serverInvocation: { invocation: 2, timestamp: 2600 },
    });
    manager.processIncomingMessages();
    const skew2 = manager.getServerSimulationDiagnostics().clockSkewMs;

    expect(skew2).toBeGreaterThan(skew1);
    expect(skew2).toBeLessThan(600);
  });

  it("propagates invocation and timestamp from server script runner", async () => {
    const schema = ServerScripts(
      Type.Object({
        count: Type.Number(),
      }),
      [
        {
          id: "loop",
          code: `
export default function loop(ctx) {
  return { ...ctx.state, count: (ctx.state.count ?? 0) + 1 };
}
`,
          on: { stateChange: true },
        },
      ]
    );

    let now = 1000;
    let state: TestState = { count: 0 };
    const invocations: { invocation?: number; timestamp?: number }[] = [];

    const manager = new ServerScriptManager({
      getQuickJSModule: () => Promise.resolve(createUnsafeQuickJSModule()),
      environment: { ...testEnvironment, clock: () => now },
      getState: () => state,
      getSharedConnectionData: () => ({}),
      setSharedConnectionData: () => {},
      drainInputQueue: () => [],
      applyState: async ({ nextState, invocation, timestamp }) => {
        state = nextState;
        invocations.push({ invocation, timestamp });
      },
      log: (_entry) => {},
    });

    await manager.updateSchema(schema);
    manager.setIsHot(true);

    manager.handleStateChange([
      { op: "replace", path: ["count"], value: 1 } as ExtendedPatch,
    ]);

    await new Promise((resolve) => setTimeout(resolve, 0));

    expect(invocations[0]).toEqual({ invocation: 1, timestamp: now });
    expect(state.count).toBe(1);

    now = 2000;
    manager.handleStateChange([
      { op: "replace", path: ["count"], value: 2 } as ExtendedPatch,
    ]);

    await new Promise((resolve) => setTimeout(resolve, 0));

    expect(invocations[1]).toEqual({ invocation: 2, timestamp: now });
    expect(state.count).toBe(2);
  });

  it("skips client simulation when server invocation is already current or newer", () => {
    const manager = new MultiplayerStateManager<TestState>({ count: 0 }, {});
    manager.connect();

    // Seed a real server invocation
    manager.incomingMessages.push({
      type: "acceptPatch",
      id: "server-1",
      patches: [],
      serverInvocation: { invocation: 2, timestamp: 1200 },
    });
    manager.processIncomingMessages();

    // Prediction for an older or current invocation should be ignored
    manager.applyServerPrediction({
      nextState: { count: 99 },
      invocation: 2,
      timestamp: 1300,
    });

    expect(manager.pendingCommits.length).toBe(0);

    // Prediction for a newer invocation should be accepted
    manager.applyServerPrediction({
      nextState: { count: 1 },
      invocation: 3,
      timestamp: 1400,
    });

    expect(manager.pendingCommits.length).toBe(1);
    expect(manager.getLatestPrediction?.()?.serverInvocation?.invocation).toBe(
      3
    );
  });

  it("starts interval simulation from latest known server invocation", async () => {
    const schema = ServerScripts(
      Type.Object({
        count: Type.Number({ default: 0 }),
      }),
      [
        {
          id: "loop",
          code: `
export default function loop(ctx) {
  return { ...ctx.state, count: (ctx.state.count ?? 0) + 1 };
}
`,
          on: { interval: "50ms" },
        },
      ]
    );

    let state: TestState = { count: 0 };
    const invocations: number[] = [];

    const manager = new ServerScriptManager({
      getQuickJSModule: () => Promise.resolve(createUnsafeQuickJSModule()),
      environment: { ...testEnvironment, clock: () => 0 },
      getState: () => state,
      getSharedConnectionData: () => ({}),
      setSharedConnectionData: () => {},
      drainInputQueue: () => [],
      getInitialInvocation: () => 10,
      applyState: async ({ nextState, invocation }) => {
        state = nextState;
        if (invocation !== undefined) invocations.push(invocation);
      },
      log: (_entry) => {},
    });

    await manager.updateSchema(schema);
    manager.setIsHot(true);

    await Bun.sleep(20);

    expect(invocations[0]).toBe(11);
  });

  it("updates simulator invocation when tracker advances after creation", async () => {
    const schema = ServerScripts(
      Type.Object({
        count: Type.Number({ default: 0 }),
      }),
      [
        {
          id: "loop",
          code: `
export default function loop(ctx) {
  return { ...ctx.state, count: (ctx.state.count ?? 0) + 1 };
}
`,
          on: { interval: "50ms" },
        },
      ]
    );

    let state: TestState = { count: 0 };
    let seed = 0;
    const invocations: number[] = [];

    const manager = new ServerScriptManager({
      getQuickJSModule: () => Promise.resolve(createUnsafeQuickJSModule()),
      environment: { ...testEnvironment, clock: () => 0 },
      getState: () => state,
      getSharedConnectionData: () => ({}),
      setSharedConnectionData: () => {},
      drainInputQueue: () => [],
      getInitialInvocation: () => seed,
      applyState: async ({ nextState, invocation }) => {
        state = nextState;
        if (invocation !== undefined) invocations.push(invocation);
      },
      log: (_entry) => {},
    });

    await manager.updateSchema(schema);
    manager.setIsHot(true);

    // Advance tracker after creation but before next tick
    seed = 20;

    await Bun.sleep(60);

    expect(invocations[0]).toBeGreaterThanOrEqual(20);
  });

  it("runs interval server script at predicted server times without falling behind", async () => {
    const schema = ServerScripts(
      Type.Object({
        tickAt: Type.Number({ default: 0 }),
        count: Type.Number({ default: 0 }),
      }),
      [
        {
          id: "loop",
          code: `
export default function loop(ctx) {
  return { ...ctx.state, tickAt: ctx.now, count: (ctx.state.count ?? 0) + 1 };
}
`,
          on: { interval: "50ms" },
        },
      ]
    );

    const start = performance.now();
    let serverAdvance = 0;
    const getServerNow = () =>
      1_000 + (performance.now() - start) + serverAdvance;
    let state: { tickAt: number; count: number } = { tickAt: 0, count: 0 };
    const log: Array<{
      invocation?: number;
      timestamp?: number;
      tickAt: number;
    }> = [];

    const manager = new ServerScriptManager({
      getQuickJSModule: () => Promise.resolve(createUnsafeQuickJSModule()),
      environment: { ...testEnvironment, clock: () => getServerNow() },
      getState: () => state,
      getSharedConnectionData: () => ({}),
      setSharedConnectionData: () => {},
      drainInputQueue: () => [],
      applyState: async ({ nextState, invocation, timestamp }) => {
        state = nextState;
        log.push({ invocation, timestamp, tickAt: nextState.tickAt });
      },
      log: (_entry) => {},
    });

    await manager.updateSchema(schema);
    manager.setIsHot(true);

    const waitForInvocationAtLeast = async (
      count: number,
      timeoutMs: number = 300
    ) => {
      const started = performance.now();
      while (
        (log.at(-1)?.invocation ?? 0) < count &&
        performance.now() - started < timeoutMs
      ) {
        await Bun.sleep(10);
      }
    };

    // Let a few ticks happen (with some wiggle room for slow CI start-up)
    await waitForInvocationAtLeast(3);
    const first = log.at(-1);
    expect(first?.invocation ?? 0).toBeGreaterThanOrEqual(3);

    const beforeJump = first?.invocation ?? 0;

    // Jump the server clock forward; the runner should catch up quickly to the predicted server time
    serverAdvance += 400;
    await waitForInvocationAtLeast(beforeJump + 3, 200);

    const last = log.at(-1);
    expect((last?.invocation ?? 0) - beforeJump).toBeGreaterThanOrEqual(3);

    // All invocations should be non-decreasing and timestamps close to reported tickAt
    for (let i = 1; i < log.length; i += 1) {
      expect(
        (log[i].invocation ?? 0) - (log[i - 1].invocation ?? 0)
      ).toBeGreaterThanOrEqual(0);
      if (log[i].timestamp !== undefined) {
        expect(Math.abs(log[i].timestamp! - log[i].tickAt)).toBeLessThanOrEqual(
          80
        );
      }
    }
  });
});