import { describe, test, expect, mock, afterEach } from "bun:test";
import type { GatewayConfig } from "../config.js";
import { initSigningKey } from "../auth/token-service.js";

const TEST_SIGNING_KEY = Buffer.from("test-signing-key-at-least-32-bytes-long");
initSigningKey(TEST_SIGNING_KEY);

type FetchFn = (
  input: string | URL | Request,
  init?: RequestInit,
) => Promise<Response>;
let fetchMock: ReturnType<typeof mock<FetchFn>> = mock(
  async () => new Response(),
);

mock.module("../fetch.js", () => ({
  fetchImpl: (...args: Parameters<FetchFn>) => fetchMock(...args),
}));

const {
  createOAuthCallbackHandler,
  _resetConsumedStates,
  MAX_CONSUMED_STATES,
} = await import("../http/routes/oauth-callback.js");

// 32-char hex string matching the format generated by the runtime (16 random bytes)
const VALID_STATE = "a1b2c3d4e5f6a7b8c9d0e1f2a3b4c5d6";
const VALID_STATE_2 = "f1e2d3c4b5a6f7e8d9c0b1a2f3e4d5c6";

function makeConfig(overrides: Partial<GatewayConfig> = {}): GatewayConfig {
  const merged: GatewayConfig = {
    assistantRuntimeBaseUrl: "http://localhost:7821",
    routingEntries: [],
    defaultAssistantId: undefined,
    unmappedPolicy: "reject",
    port: 7830,
    runtimeProxyRequireAuth: true,
    shutdownDrainMs: 5000,
    runtimeTimeoutMs: 30000,
    runtimeMaxRetries: 2,
    runtimeInitialBackoffMs: 500,
    maxWebhookPayloadBytes: 1048576,
    logFile: { dir: undefined, retentionDays: 30 },
    maxAttachmentBytes: {
      telegram: 50 * 1024 * 1024,
      slack: 100 * 1024 * 1024,
      whatsapp: 16 * 1024 * 1024,
      default: 50 * 1024 * 1024,
    },
    maxAttachmentConcurrency: 3,
    gatewayInternalBaseUrl: "http://127.0.0.1:7830",
    trustProxy: false,
    ...overrides,
  };
  return merged;
}

describe("OAuth callback handler", () => {
  afterEach(() => {
    fetchMock = mock(async () => new Response());
    _resetConsumedStates();
  });

  test("forwards valid state+code to runtime and returns success page", async () => {
    fetchMock = mock(async () => Response.json({ ok: true }, { status: 200 }));

    const handler = createOAuthCallbackHandler(makeConfig());
    const req = new Request(
      `http://localhost:7830/webhooks/oauth/callback?state=${VALID_STATE}&code=authcode456`,
      { method: "GET" },
    );

    const res = await handler(req);
    expect(res.status).toBe(200);

    const body = await res.text();
    expect(body).toContain("Authorization Successful");
    expect(body).toContain("close this tab");

    // Verify fetch was called with the runtime internal endpoint
    const calledUrl = (fetchMock.mock.calls[0] as unknown[])[0] as string;
    expect(calledUrl).toBe("http://localhost:7821/v1/internal/oauth/callback");

    const calledInit = (fetchMock.mock.calls[0] as unknown[])[1] as RequestInit;
    const sentBody = JSON.parse(calledInit.body as string);
    expect(sentBody.state).toBe(VALID_STATE);
    expect(sentBody.code).toBe("authcode456");
    expect(sentBody.error).toBeUndefined();

    const headers = calledInit.headers as Record<string, string>;
    expect(headers["Authorization"]).toMatch(/^Bearer ey/);
  });

  test("missing state parameter returns 400 error page", async () => {
    const handler = createOAuthCallbackHandler(makeConfig());
    const req = new Request(
      "http://localhost:7830/webhooks/oauth/callback?code=authcode456",
      { method: "GET" },
    );

    const res = await handler(req);
    expect(res.status).toBe(400);

    const body = await res.text();
    expect(body).toContain("Authorization Failed");
    expect(body).toContain("Missing state parameter");
  });

  test("state shorter than 32 chars returns 400", async () => {
    const handler = createOAuthCallbackHandler(makeConfig());
    const req = new Request(
      "http://localhost:7830/webhooks/oauth/callback?state=abc123&code=authcode456",
      { method: "GET" },
    );

    const res = await handler(req);
    expect(res.status).toBe(400);

    const body = await res.text();
    expect(body).toContain("Invalid state parameter");
    expect(fetchMock.mock.calls.length).toBe(0);
  });

  test("replayed state returns 400", async () => {
    fetchMock = mock(async () => Response.json({ ok: true }, { status: 200 }));

    const handler = createOAuthCallbackHandler(makeConfig());

    // First request succeeds
    const req1 = new Request(
      `http://localhost:7830/webhooks/oauth/callback?state=${VALID_STATE}&code=code1`,
      { method: "GET" },
    );
    const res1 = await handler(req1);
    expect(res1.status).toBe(200);

    // Replay with the same state is blocked
    const req2 = new Request(
      `http://localhost:7830/webhooks/oauth/callback?state=${VALID_STATE}&code=code2`,
      { method: "GET" },
    );
    const res2 = await handler(req2);
    expect(res2.status).toBe(400);

    const body = await res2.text();
    expect(body).toContain("State parameter already used");
    // Only the first request should have been forwarded
    expect(fetchMock.mock.calls.length).toBe(1);
  });

  test("different states are not blocked by replay check", async () => {
    fetchMock = mock(async () => Response.json({ ok: true }, { status: 200 }));

    const handler = createOAuthCallbackHandler(makeConfig());

    const req1 = new Request(
      `http://localhost:7830/webhooks/oauth/callback?state=${VALID_STATE}&code=code1`,
      { method: "GET" },
    );
    await handler(req1);

    const req2 = new Request(
      `http://localhost:7830/webhooks/oauth/callback?state=${VALID_STATE_2}&code=code2`,
      { method: "GET" },
    );
    const res2 = await handler(req2);
    expect(res2.status).toBe(200);
    expect(fetchMock.mock.calls.length).toBe(2);
  });

  test("error param is forwarded to runtime", async () => {
    fetchMock = mock(async () => Response.json({ ok: true }, { status: 200 }));

    const handler = createOAuthCallbackHandler(makeConfig());
    const req = new Request(
      `http://localhost:7830/webhooks/oauth/callback?state=${VALID_STATE}&error=access_denied`,
      { method: "GET" },
    );

    const res = await handler(req);
    expect(res.status).toBe(200);

    const calledInit = (fetchMock.mock.calls[0] as unknown[])[1] as RequestInit;
    const sentBody = JSON.parse(calledInit.body as string);
    expect(sentBody.state).toBe(VALID_STATE);
    expect(sentBody.error).toBe("access_denied");
    expect(sentBody.code).toBeUndefined();
  });

  test("runtime returning 404 (deterministic rejection) keeps state consumed", async () => {
    fetchMock = mock(async () =>
      Response.json({ error: "Unknown state" }, { status: 404 }),
    );

    const handler = createOAuthCallbackHandler(makeConfig());
    const req1 = new Request(
      `http://localhost:7830/webhooks/oauth/callback?state=${VALID_STATE}&code=code`,
      { method: "GET" },
    );

    const res1 = await handler(req1);
    expect(res1.status).toBe(400);
    expect(await res1.text()).toContain("Authorization Failed");

    // State stays consumed — replay is blocked even after a 4xx rejection
    const req2 = new Request(
      `http://localhost:7830/webhooks/oauth/callback?state=${VALID_STATE}&code=code`,
      { method: "GET" },
    );
    const res2 = await handler(req2);
    expect(res2.status).toBe(400);
    expect(await res2.text()).toContain("State parameter already used");
    // Only the first request should have been forwarded to the runtime
    expect(fetchMock.mock.calls.length).toBe(1);
  });

  test("runtime returning 5xx allows retry", async () => {
    let callCount = 0;
    fetchMock = mock(async () => {
      callCount++;
      if (callCount === 1) {
        return Response.json({ error: "Internal error" }, { status: 500 });
      }
      return Response.json({ ok: true }, { status: 200 });
    });

    const handler = createOAuthCallbackHandler(makeConfig());
    const req1 = new Request(
      `http://localhost:7830/webhooks/oauth/callback?state=${VALID_STATE}&code=code`,
      { method: "GET" },
    );

    const res1 = await handler(req1);
    expect(res1.status).toBe(400);
    expect(await res1.text()).toContain("Authorization Failed");

    // State should be rolled back on 5xx so a retry succeeds
    const req2 = new Request(
      `http://localhost:7830/webhooks/oauth/callback?state=${VALID_STATE}&code=code`,
      { method: "GET" },
    );
    const res2 = await handler(req2);
    expect(res2.status).toBe(200);
    expect(callCount).toBe(2);
  });

  test("runtime unreachable returns 502 error page and allows retry", async () => {
    let callCount = 0;
    fetchMock = mock(async () => {
      callCount++;
      if (callCount === 1) {
        throw new Error("Connection refused");
      }
      return Response.json({ ok: true }, { status: 200 });
    });

    const handler = createOAuthCallbackHandler(makeConfig());
    const req1 = new Request(
      `http://localhost:7830/webhooks/oauth/callback?state=${VALID_STATE}&code=code`,
      { method: "GET" },
    );

    const res1 = await handler(req1);
    expect(res1.status).toBe(502);
    expect(await res1.text()).toContain("Authorization Failed");

    // State should be rolled back so a retry succeeds
    const req2 = new Request(
      `http://localhost:7830/webhooks/oauth/callback?state=${VALID_STATE}&code=code`,
      { method: "GET" },
    );
    const res2 = await handler(req2);
    expect(res2.status).toBe(200);
    expect(callCount).toBe(2);
  });

  test("evicts oldest entry when consumed-state map is at capacity", async () => {
    fetchMock = mock(async () => Response.json({ ok: true }, { status: 200 }));

    const handler = createOAuthCallbackHandler(makeConfig());

    // Fill the map to capacity
    for (let i = 0; i < MAX_CONSUMED_STATES; i++) {
      const state = `a1b2c3d4e5f6a7b8c9d0e1f2a3b4${i.toString().padStart(4, "0")}`;
      const req = new Request(
        `http://localhost:7830/webhooks/oauth/callback?state=${state}&code=code`,
        { method: "GET" },
      );
      await handler(req);
    }

    // Next callback should succeed — LRU eviction drops the oldest entry
    const req = new Request(
      `http://localhost:7830/webhooks/oauth/callback?state=${VALID_STATE}&code=code`,
      { method: "GET" },
    );
    const res = await handler(req);
    expect(res.status).toBe(200);

    // The new state was forwarded to the runtime
    expect(fetchMock.mock.calls.length).toBe(MAX_CONSUMED_STATES + 1);
  });

  test("always sends JWT Authorization header to runtime", async () => {
    fetchMock = mock(async () => Response.json({ ok: true }, { status: 200 }));

    const handler = createOAuthCallbackHandler(makeConfig({}));
    const req = new Request(
      `http://localhost:7830/webhooks/oauth/callback?state=${VALID_STATE}&code=c1`,
      { method: "GET" },
    );

    await handler(req);

    const calledInit = (fetchMock.mock.calls[0] as unknown[])[1] as RequestInit;
    const headers = calledInit.headers as Record<string, string>;
    expect(headers["Authorization"]).toMatch(/^Bearer ey/);
  });
});