import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
import { JsonRpcTransport } from "./json-rpc-transport.js";

// ─── Helpers ─────────────────────────────────────────────────────────────────

/**
 * Creates a controllable ReadableStream that simulates stdout.
 * Returns an enqueue callback for pushing data and a close callback.
 */
function createMockStdout(): {
  stream: ReadableStream<Uint8Array>;
  push: (text: string) => void;
  close: () => void;
} {
  let ctrl!: ReadableStreamDefaultController<Uint8Array>;
  const stream = new ReadableStream<Uint8Array>({
    start(c) {
      ctrl = c;
    },
  });
  const encoder = new TextEncoder();
  return {
    stream,
    push: (text: string) => ctrl.enqueue(encoder.encode(text)),
    close: () => ctrl.close(),
  };
}

/**
 * Creates a WritableStream that captures all written bytes as strings,
 * simulating stdin for the transport.
 */
function createMockStdin(): {
  stream: WritableStream<Uint8Array>;
  written: string[];
} {
  const decoder = new TextDecoder();
  const written: string[] = [];
  const stream = new WritableStream<Uint8Array>({
    write(chunk) {
      written.push(decoder.decode(chunk));
    },
  });
  return { stream, written };
}

/**
 * Flushes microtask queue so the transport's background readStdout loop
 * processes newly pushed data.
 */
async function flushAsync(): Promise<void> {
  await new Promise((r) => setTimeout(r, 10));
}

// =============================================================================
// Tests
// =============================================================================

describe("JsonRpcTransport", () => {
  let transport: JsonRpcTransport;
  let stdout: ReturnType<typeof createMockStdout>;
  let stdin: ReturnType<typeof createMockStdin>;

  beforeEach(() => {
    stdout = createMockStdout();
    stdin = createMockStdin();
    transport = new JsonRpcTransport(stdin.stream, stdout.stream);
  });

  afterEach(() => {
    // Close stdout to stop the background reader loop and avoid dangling promises
    try {
      stdout.close();
    } catch {
      // already closed
    }
  });

  // ─── Connection status ───────────────────────────────────────────────────

  describe("isConnected", () => {
    it("returns true immediately after construction", () => {
      expect(transport.isConnected()).toBe(true);
    });

    it("returns false after stdout stream closes", async () => {
      stdout.close();
      await flushAsync();
      expect(transport.isConnected()).toBe(false);
    });
  });

  // ─── Sending messages ────────────────────────────────────────────────────

  describe("call", () => {
    it("serializes a JSON-RPC request with auto-incrementing id", async () => {
      // Don't await — the call will block waiting for a response
      const promise1 = transport.call("method1", { key: "val1" });
      const promise2 = transport.call("method2", { key: "val2" });
      await flushAsync();

      // Both requests should be written to stdin as newline-delimited JSON
      expect(stdin.written.length).toBe(2);

      const req1 = JSON.parse(stdin.written[0].replace("\n", ""));
      expect(req1).toEqual({ method: "method1", id: 1, params: { key: "val1" } });

      const req2 = JSON.parse(stdin.written[1].replace("\n", ""));
      expect(req2).toEqual({ method: "method2", id: 2, params: { key: "val2" } });

      // Simulate responses so the promises don't hang
      stdout.push(JSON.stringify({ id: 1, result: "ok1" }) + "\n");
      stdout.push(JSON.stringify({ id: 2, result: "ok2" }) + "\n");
      await flushAsync();

      expect(await promise1).toBe("ok1");
      expect(await promise2).toBe("ok2");
    });

    it("rejects when the response contains an error", async () => {
      const promise = transport.call("failMethod");
      // Set up the rejection handler before the error arrives to avoid unhandled rejection warning
      const expectation = expect(promise).rejects.toThrow("boom");
      await flushAsync();

      stdout.push(JSON.stringify({ id: 1, error: { code: -1, message: "boom" } }) + "\n");
      await flushAsync();

      await expectation;
    });

    it("rejects all pending calls when transport closes", async () => {
      const promise = transport.call("neverResponds");
      // Set up the rejection handler before closing to avoid unhandled rejection warning
      const expectation = expect(promise).rejects.toThrow("Transport closed");
      await flushAsync();

      stdout.close();
      await flushAsync();

      await expectation;
    });

    it("rejects immediately when transport is already closed", async () => {
      stdout.close();
      await flushAsync();

      await expect(transport.call("late")).rejects.toThrow("Transport closed");
    });
  });

  describe("notify", () => {
    it("sends a notification without an id field", async () => {
      await transport.notify("someEvent", { data: 42 });
      await flushAsync();

      expect(stdin.written.length).toBe(1);
      const msg = JSON.parse(stdin.written[0].replace("\n", ""));
      expect(msg).toEqual({ method: "someEvent", params: { data: 42 } });
      expect(msg).not.toHaveProperty("id");
    });
  });

  describe("respond", () => {
    it("sends a response with the given id and result", async () => {
      await transport.respond(7, { status: "approved" });
      await flushAsync();

      expect(stdin.written.length).toBe(1);
      const msg = JSON.parse(stdin.written[0].replace("\n", ""));
      expect(msg).toEqual({ id: 7, result: { status: "approved" } });
    });
  });

  // ─── Receiving messages ──────────────────────────────────────────────────

  describe("onNotification", () => {
    it("dispatches server-initiated notifications (no id)", async () => {
      const handler = vi.fn();
      transport.onNotification(handler);

      stdout.push(JSON.stringify({ method: "progress", params: { pct: 50 } }) + "\n");
      await flushAsync();

      expect(handler).toHaveBeenCalledWith("progress", { pct: 50 });
    });

    it("does not trigger notification handler for responses", async () => {
      const handler = vi.fn();
      transport.onNotification(handler);

      // Start a call so there's a pending entry
      const p = transport.call("test");
      await flushAsync();

      stdout.push(JSON.stringify({ id: 1, result: "done" }) + "\n");
      await flushAsync();
      await p;

      expect(handler).not.toHaveBeenCalled();
    });
  });

  describe("onRequest", () => {
    it("dispatches server-initiated requests (has id + method)", async () => {
      const handler = vi.fn();
      transport.onRequest(handler);

      stdout.push(JSON.stringify({ method: "approval", id: 99, params: { tool: "bash" } }) + "\n");
      await flushAsync();

      expect(handler).toHaveBeenCalledWith("approval", 99, { tool: "bash" });
    });
  });

  // ─── Raw callbacks ────────────────────────────────────────────────────────

  describe("raw callbacks", () => {
    it("fires onRawIncoming before dispatching", async () => {
      const rawCb = vi.fn();
      transport.onRawIncoming(rawCb);

      const line = JSON.stringify({ method: "event", params: {} });
      stdout.push(line + "\n");
      await flushAsync();

      expect(rawCb).toHaveBeenCalledWith(line);
    });

    it("fires onRawOutgoing when sending", async () => {
      const rawCb = vi.fn();
      transport.onRawOutgoing(rawCb);

      await transport.notify("ping", {});
      await flushAsync();

      expect(rawCb).toHaveBeenCalledTimes(1);
      const outgoing = rawCb.mock.calls[0][0] as string;
      expect(outgoing).toContain("ping");
    });
  });

  // ─── Buffer handling ──────────────────────────────────────────────────────

  describe("buffer handling", () => {
    it("handles messages split across multiple chunks", async () => {
      const handler = vi.fn();
      transport.onNotification(handler);

      // Push a single notification in two chunks
      const full = JSON.stringify({ method: "split", params: { ok: true } });
      const mid = Math.floor(full.length / 2);
      stdout.push(full.slice(0, mid));
      await flushAsync();
      expect(handler).not.toHaveBeenCalled();

      stdout.push(full.slice(mid) + "\n");
      await flushAsync();
      expect(handler).toHaveBeenCalledWith("split", { ok: true });
    });

    it("handles multiple messages in a single chunk", async () => {
      const handler = vi.fn();
      transport.onNotification(handler);

      const msg1 = JSON.stringify({ method: "a", params: {} });
      const msg2 = JSON.stringify({ method: "b", params: {} });
      stdout.push(msg1 + "\n" + msg2 + "\n");
      await flushAsync();

      expect(handler).toHaveBeenCalledTimes(2);
      expect(handler).toHaveBeenCalledWith("a", {});
      expect(handler).toHaveBeenCalledWith("b", {});
    });

    it("ignores blank lines and non-JSON lines gracefully", async () => {
      const handler = vi.fn();
      transport.onNotification(handler);

      stdout.push("\n\n   \nnot-json-at-all\n");
      await flushAsync();

      // No crash, no calls
      expect(handler).not.toHaveBeenCalled();
    });
  });

  // ─── Bun-style stdin (object with write method) ───────────────────────────

  describe("Bun-style stdin object", () => {
    it("wraps a { write() } object into a WritableStream", async () => {
      const chunks: Uint8Array[] = [];
      const bunStdin = { write: (data: Uint8Array) => { chunks.push(data); return data.length; } };

      const stdout2 = createMockStdout();
      const transport2 = new JsonRpcTransport(bunStdin, stdout2.stream);

      await transport2.notify("hello");
      await flushAsync();

      expect(chunks.length).toBe(1);
      const decoded = new TextDecoder().decode(chunks[0]);
      expect(JSON.parse(decoded.replace("\n", ""))).toEqual({ method: "hello", params: {} });

      stdout2.close();
    });
  });

  // ─── Default empty params ─────────────────────────────────────────────────

  describe("default empty params", () => {
    it("call uses empty object as default params", async () => {
      const promise = transport.call("noParams");
      await flushAsync();

      expect(stdin.written.length).toBe(1);
      const req = JSON.parse(stdin.written[0].replace("\n", ""));
      expect(req.params).toEqual({});

      // Send response so promise resolves
      stdout.push(JSON.stringify({ id: 1, result: "ok" }) + "\n");
      await flushAsync();
      await promise;
    });

    it("notify uses empty object as default params", async () => {
      await transport.notify("noParamsEvent");
      await flushAsync();

      const msg = JSON.parse(stdin.written[0].replace("\n", ""));
      expect(msg.params).toEqual({});
    });
  });

  // ─── Missing params in incoming messages ──────────────────────────────────

  describe("missing params in incoming messages", () => {
    it("defaults to empty object for notifications without params", async () => {
      const handler = vi.fn();
      transport.onNotification(handler);

      stdout.push(JSON.stringify({ method: "bare" }) + "\n");
      await flushAsync();

      expect(handler).toHaveBeenCalledWith("bare", {});
    });

    it("defaults to empty object for requests without params", async () => {
      const handler = vi.fn();
      transport.onRequest(handler);

      stdout.push(JSON.stringify({ method: "ask", id: 42 }) + "\n");
      await flushAsync();

      expect(handler).toHaveBeenCalledWith("ask", 42, {});
    });
  });

  // ─── Orphaned response IDs ────────────────────────────────────────────────

  describe("orphaned response IDs", () => {
    it("ignores response for unknown/already-resolved ID without crashing", async () => {
      const handler = vi.fn();
      transport.onNotification(handler);

      // Send a response with an ID that no call is waiting for
      stdout.push(JSON.stringify({ id: 999, result: "orphan" }) + "\n");
      await flushAsync();

      // Should not crash and should not be dispatched as a notification
      expect(handler).not.toHaveBeenCalled();
    });
  });

  // ─── Operations after transport close ─────────────────────────────────────

  describe("operations after close", () => {
    it("rejects calls after transport is closed", async () => {
      stdout.close();
      await flushAsync();

      await expect(transport.call("late")).rejects.toThrow("Transport closed");
    });

    it("rejects notify after transport is closed", async () => {
      stdout.close();
      await flushAsync();

      await expect(transport.notify("lateEvent")).rejects.toThrow("Transport closed");
    });

    it("rejects respond after transport is closed", async () => {
      stdout.close();
      await flushAsync();

      await expect(transport.respond(1, "late")).rejects.toThrow("Transport closed");
    });
  });

  // ─── ID isolation between instances ───────────────────────────────────────

  describe("ID isolation between instances", () => {
    it("each transport instance starts IDs from 1", async () => {
      // First transport
      const p1 = transport.call("method1");
      await flushAsync();
      const req1 = JSON.parse(stdin.written[0].replace("\n", ""));
      expect(req1.id).toBe(1);

      // Resolve the call
      stdout.push(JSON.stringify({ id: 1, result: "ok" }) + "\n");
      await flushAsync();
      await p1;

      // Second transport with its own streams
      const stdout2 = createMockStdout();
      const stdin2 = createMockStdin();
      const transport2 = new JsonRpcTransport(stdin2.stream, stdout2.stream);

      const p2 = transport2.call("method2");
      await flushAsync();

      // Second transport should also start at id 1
      const req2 = JSON.parse(stdin2.written[0].replace("\n", ""));
      expect(req2.id).toBe(1);

      // Resolve and close
      stdout2.push(JSON.stringify({ id: 1, result: "ok2" }) + "\n");
      await flushAsync();
      await p2;
      stdout2.close();
    });
  });
});
