import { describe, expect, it, vi } from "vitest";
import type { WebSocket } from "ws";
import { UserStreamMux, startServerPing, type StreamContext } from "../src/stream.js";
import type { ClientMessage, ServerMessage, Session } from "../src/types.js";
import { flushMicrotasks } from "./harness/async.js";
import { FakeWebSocket, makeSession } from "./harness/stream-test-primitives.js";
import { messagesOfType, requireMessageOfType } from "./harness/ws-harness.js";

interface Harness {
  mux: UserStreamMux;
  sessionCallbacks: Map<string, (msg: ServerMessage) => void>;
  unsubscribeCalls: string[];
  handleClientMessage: ReturnType<typeof vi.fn>;
  getCatchUp: ReturnType<typeof vi.fn>;
}

function makeHarness(options?: {
  sessions?: Session[];
  catchUpBySession?: Record<
    string,
    {
      events: ServerMessage[];
      currentSeq: number;
      catchUpComplete: boolean;
    }
  >;
  ensureSessionContextWindow?: (session: Session) => Session;
}): Harness {
  const sessions = options?.sessions ?? [makeSession("s1")];
  const sessionsById = new Map(sessions.map((session) => [session.id, session]));
  const sessionCallbacks = new Map<string, (msg: ServerMessage) => void>();
  const unsubscribeCalls: string[] = [];

  const getCurrentSeq = vi.fn((sessionId: string) => {
    const fixture = options?.catchUpBySession?.[sessionId];
    return fixture?.currentSeq ?? 0;
  });

  const getCatchUp = vi.fn((sessionId: string, _sinceSeq: number) => {
    const session = sessionsById.get(sessionId);
    if (!session) return null;

    const fixture = options?.catchUpBySession?.[sessionId];
    if (!fixture) {
      return { events: [], currentSeq: 0, session, catchUpComplete: true };
    }
    return {
      events: fixture.events,
      currentSeq: fixture.currentSeq,
      session,
      catchUpComplete: fixture.catchUpComplete,
    };
  });

  const handleClientMessage = vi.fn(
    async (_session: Session, _msg: ClientMessage, send: (msg: ServerMessage) => void) => {
      send({ type: "agent_start" });
    },
  );

  const ctx: StreamContext = {
    sessions: {
      startSession: vi.fn(async (sessionId: string) => {
        const session = sessionsById.get(sessionId);
        if (!session) throw new Error(`Session not found: ${sessionId}`);
        return session;
      }),
      subscribe: vi.fn((sessionId: string, cb: (msg: ServerMessage) => void) => {
        sessionCallbacks.set(sessionId, cb);
        return () => {
          unsubscribeCalls.push(sessionId);
          if (sessionCallbacks.get(sessionId) === cb) {
            sessionCallbacks.delete(sessionId);
          }
        };
      }),
      getActiveSession: vi.fn((sessionId: string) => sessionsById.get(sessionId)),
      getCurrentSeq,
      getCatchUp,
      getPendingAskMessage: () => undefined,
    } as unknown as StreamContext["sessions"],
    storage: {
      getSession: vi.fn((sessionId: string) => sessionsById.get(sessionId)),
      getOwnerName: vi.fn(() => "test-host"),
    } as unknown as StreamContext["storage"],
    gate: {
      getPendingForUser: vi.fn(() => []),
      resolveDecision: vi.fn(() => true),
    } as unknown as StreamContext["gate"],
    ensureSessionContextWindow:
      options?.ensureSessionContextWindow ?? ((session: Session) => session),
    resolveWorkspaceForSession: () => undefined,
    handleClientMessage,
    trackConnection: vi.fn(),
    untrackConnection: vi.fn(),
  };

  const mux = new UserStreamMux(ctx);

  return {
    mux,
    sessionCallbacks,
    unsubscribeCalls,
    handleClientMessage,
    getCatchUp,
  };
}

async function flushQueue(): Promise<void> {
  await flushMicrotasks(3);
}

describe("/stream websocket behavior", () => {
  it("streams active and background session events on one socket", async () => {
    const harness = makeHarness({ sessions: [makeSession("s1"), makeSession("s2")] });
    const ws = new FakeWebSocket();

    await harness.mux.handleWebSocket(ws as unknown as WebSocket);

    ws.emitClientMessage({
      type: "subscribe",
      sessionId: "s1",
      level: "full",
      requestId: "sub-1",
    });
    await flushQueue();
    ws.emitClientMessage({
      type: "subscribe",
      sessionId: "s2",
      level: "notifications",
      requestId: "sub-2",
    });
    await flushQueue();

    const base = ws.sent.length;

    harness.sessionCallbacks.get("s1")?.({ type: "text_delta", delta: "active-delta" });
    harness.sessionCallbacks.get("s1")?.({ type: "agent_start" });
    harness.sessionCallbacks.get("s2")?.({ type: "text_delta", delta: "bg-delta" });
    harness.sessionCallbacks.get("s2")?.({ type: "agent_end" });

    const delivered = ws.sent.slice(base);
    expect(delivered).toContainEqual({
      type: "text_delta",
      delta: "active-delta",
      sessionId: "s1",
    });
    expect(delivered).toContainEqual({ type: "agent_start", sessionId: "s1" });
    expect(delivered).toContainEqual({ type: "agent_end", sessionId: "s2" });
    expect(
      delivered.find((msg) => msg.type === "text_delta" && msg.sessionId === "s2"),
    ).toBeUndefined();
  });

  it("normalizes stale context window on streamed state events", async () => {
    const staleSession: Session = {
      ...makeSession("s1"),
      model: "gpt-5.3-codex/gpt-5.3-codex",
      contextWindow: 200000,
    };

    const harness = makeHarness({
      sessions: [staleSession],
      ensureSessionContextWindow: (session: Session) => {
        if (session.model?.includes("gpt-5.3-codex") && session.contextWindow === 200000) {
          return { ...session, contextWindow: 272000 };
        }
        return session;
      },
    });
    const ws = new FakeWebSocket();

    await harness.mux.handleWebSocket(ws as unknown as WebSocket);

    ws.emitClientMessage({
      type: "subscribe",
      sessionId: "s1",
      level: "full",
      requestId: "sub-1",
    });
    await flushQueue();

    const callback = harness.sessionCallbacks.get("s1");
    expect(callback).toBeTruthy();
    callback?.({ type: "state", session: staleSession });

    const latestState = [...ws.sent].reverse().find((msg) => msg.type === "state");
    expect(latestState?.type).toBe("state");
    if (latestState?.type === "state") {
      expect(latestState.session.contextWindow).toBe(272000);
    }
  });

  it("deduplicates same-level subscribe without re-subscribing", async () => {
    const harness = makeHarness();
    const ws = new FakeWebSocket();

    await harness.mux.handleWebSocket(ws as unknown as WebSocket);

    ws.emitClientMessage({
      type: "subscribe",
      sessionId: "s1",
      level: "notifications",
      requestId: "sub-1",
    });
    await flushQueue();

    // Second subscribe at same level should be deduplicated — no unsubscribe/resubscribe.
    ws.emitClientMessage({
      type: "subscribe",
      sessionId: "s1",
      level: "notifications",
      requestId: "sub-2",
    });
    await flushQueue();

    // Dedup: no clearSubscription called for the second subscribe.
    expect(harness.unsubscribeCalls).toEqual([]);

    const subscribeResults = messagesOfType(ws.sent, "command_result").filter(
      (msg) => msg.command === "subscribe",
    );
    expect(subscribeResults).toHaveLength(2);
    expect(subscribeResults.every((msg) => msg.success)).toBe(true);
    // Second result should be marked as deduplicated.
    expect(subscribeResults[1].data?.deduplicated).toBe(true);

    // Unsubscribe still works normally after dedup.
    ws.emitClientMessage({ type: "unsubscribe", sessionId: "s1", requestId: "unsub-1" });
    await flushQueue();
    ws.emitClientMessage({ type: "unsubscribe", sessionId: "s1", requestId: "unsub-2" });
    await flushQueue();

    const unsubscribeResults = messagesOfType(ws.sent, "command_result").filter(
      (msg) => msg.command === "unsubscribe",
    );

    expect(unsubscribeResults).toHaveLength(2);
    expect(unsubscribeResults.every((msg) => msg.success)).toBe(true);
  });

  it("allows multiple full subscriptions without downgrading", async () => {
    const harness = makeHarness({ sessions: [makeSession("s1"), makeSession("s2")] });
    const ws = new FakeWebSocket();

    await harness.mux.handleWebSocket(ws as unknown as WebSocket);

    ws.emitClientMessage({
      type: "subscribe",
      sessionId: "s1",
      level: "full",
      requestId: "sub-1",
    });
    await flushQueue();
    ws.emitClientMessage({
      type: "subscribe",
      sessionId: "s2",
      level: "full",
      requestId: "sub-2",
    });
    await flushQueue();

    // Both sessions should accept commands — no auto-downgrade
    const beforeS1Prompt = ws.sent.length;
    ws.emitClientMessage({
      type: "prompt",
      sessionId: "s1",
      message: "should succeed",
      requestId: "prompt-s1",
    });
    await flushQueue();

    // No "not subscribed at level=full" error for s1
    const s1Errors = ws.sent
      .slice(beforeS1Prompt)
      .filter((msg) => msg.type === "error" && msg.sessionId === "s1");
    expect(s1Errors).toHaveLength(0);

    const beforeS2Prompt = ws.sent.length;
    ws.emitClientMessage({
      type: "prompt",
      sessionId: "s2",
      message: "should also succeed",
      requestId: "prompt-s2",
    });
    await flushQueue();

    const s2Errors = ws.sent
      .slice(beforeS2Prompt)
      .filter((msg) => msg.type === "error" && msg.sessionId === "s2");
    expect(s2Errors).toHaveLength(0);

    // handleClientMessage should have been called for both sessions
    expect(harness.handleClientMessage).toHaveBeenCalledTimes(2);
  });

  // Backpressure dropping was intentionally removed — it was dropping tool output
  // and thinking deltas that the client needs. All events are delivered unconditionally.

  it("reconnect + re-subscribe with sinceSeq is deterministic", async () => {
    const catchUpEvents: ServerMessage[] = [
      { type: "agent_start", seq: 6 },
      { type: "agent_end", seq: 7 },
    ];

    const harness = makeHarness({
      sessions: [makeSession("s1")],
      catchUpBySession: {
        s1: {
          events: catchUpEvents,
          currentSeq: 7,
          catchUpComplete: true,
        },
      },
    });

    const connect = async (): Promise<FakeWebSocket> => {
      const ws = new FakeWebSocket();
      await harness.mux.handleWebSocket(ws as unknown as WebSocket);
      return ws;
    };

    const ws1 = await connect();
    ws1.emitClientMessage({
      type: "subscribe",
      sessionId: "s1",
      level: "full",
      sinceSeq: 5,
      requestId: "sub-1",
    });
    await flushQueue();
    ws1.emitClose();

    const ws2 = await connect();
    ws2.emitClientMessage({
      type: "subscribe",
      sessionId: "s1",
      level: "full",
      sinceSeq: 5,
      requestId: "sub-1",
    });
    await flushQueue();

    const normalize = (messages: ServerMessage[]) =>
      messages.map((msg) => {
        if (msg.type === "command_result") {
          return {
            type: msg.type,
            command: msg.command,
            requestId: msg.requestId,
            success: msg.success,
            sessionId: msg.sessionId,
            data: msg.data,
          };
        }
        if (msg.type === "state") {
          return {
            type: msg.type,
            sessionId: msg.sessionId,
            status: msg.session.status,
          };
        }
        if (msg.type === "connected") {
          return {
            type: msg.type,
            sessionId: msg.sessionId,
            currentSeq: msg.currentSeq,
          };
        }
        return {
          type: msg.type,
          sessionId: msg.sessionId,
          seq: msg.seq,
        };
      });

    expect(normalize(ws2.sent)).toEqual(normalize(ws1.sent));

    expect(harness.getCatchUp).toHaveBeenCalledWith("s1", 5);
  });

  it("rejects inbound messages without a type field", async () => {
    const harness = makeHarness();
    const ws = new FakeWebSocket();

    await harness.mux.handleWebSocket(ws as unknown as WebSocket);

    ws.emitMessage({ requestId: "req-no-type" });
    await flushQueue();

    const error = requireMessageOfType(ws.sent, "error", (msg) =>
      msg.error.includes("Message type is required"),
    );
    expect(error.error).toContain("Message type is required");
    expect(harness.handleClientMessage).not.toHaveBeenCalled();
  });

  it("rejects unknown inbound command types without sessionId", async () => {
    const harness = makeHarness();
    const ws = new FakeWebSocket();

    await harness.mux.handleWebSocket(ws as unknown as WebSocket);

    // Unknown types without sessionId are rejected by the stream mux
    // before reaching handleClientMessage.
    ws.emitMessage({ type: "future_command_v99", requestId: "req-unknown" });
    await flushQueue();

    const error = requireMessageOfType(ws.sent, "error", (msg) =>
      msg.error.includes("sessionId is required"),
    );
    expect(error.error).toContain("future_command_v99");
    expect(harness.handleClientMessage).not.toHaveBeenCalled();
  });
});

describe("server-side ping keepalive", () => {
  it("terminates connection after 2 consecutive missed pongs", async () => {
    let terminated = false;

    const fakeWs = {
      on: vi.fn((event: string, handler: () => void) => {
        // Don't register a pong handler — simulates no pong
        if (event === "pong") {
          // Intentionally do NOT call handler — pong never fires
        }
      }),
      ping: vi.fn(() => {
        // No-op — client doesn't respond
      }),
      terminate: vi.fn(() => {
        terminated = true;
      }),
    };

    vi.useFakeTimers();
    try {
      const stop = startServerPing(fakeWs as unknown as WebSocket, "test", 100);

      // First tick: missedPongs=1, send ping — tolerated
      vi.advanceTimersByTime(100);
      expect(fakeWs.ping).toHaveBeenCalledTimes(1);
      expect(terminated).toBe(false);

      // Second tick: missedPongs=2, send ping — still tolerated (threshold is >2)
      vi.advanceTimersByTime(100);
      expect(fakeWs.ping).toHaveBeenCalledTimes(2);
      expect(terminated).toBe(false);

      // Third tick: missedPongs=3 (>2) → terminate
      vi.advanceTimersByTime(100);
      expect(terminated).toBe(true);
      expect(fakeWs.terminate).toHaveBeenCalledTimes(1);

      stop();
    } finally {
      vi.useRealTimers();
    }
  });

  it("stays alive when pong is received", async () => {
    let pongHandler: (() => void) | null = null;
    let terminated = false;

    const fakeWs = {
      on: vi.fn((event: string, handler: () => void) => {
        if (event === "pong") {
          pongHandler = handler;
        }
      }),
      ping: vi.fn(),
      terminate: vi.fn(() => {
        terminated = true;
      }),
    };

    vi.useFakeTimers();
    try {
      const stop = startServerPing(fakeWs as unknown as WebSocket, "test", 100);

      // Tick 1: send ping
      vi.advanceTimersByTime(100);
      expect(fakeWs.ping).toHaveBeenCalledTimes(1);

      // Simulate pong received
      pongHandler?.();

      // Tick 2: alive was reset → send another ping (not terminate)
      vi.advanceTimersByTime(100);
      expect(fakeWs.ping).toHaveBeenCalledTimes(2);
      expect(terminated).toBe(false);

      // Pong again
      pongHandler?.();

      // Tick 3: still alive
      vi.advanceTimersByTime(100);
      expect(fakeWs.ping).toHaveBeenCalledTimes(3);
      expect(terminated).toBe(false);

      stop();
    } finally {
      vi.useRealTimers();
    }
  });
});