import { describe, expect, it } from "vitest";
import { SignJWT, generateKeyPair } from "jose";
import { makeVerifier, rolesFromClaims } from "../server/utils/auth";

const ISS = "https://auth.kuns.dev";
const ROLES_CLAIM = "urn:zitadel:iam:org:project:roles";
const SCOPED_ROLES_CLAIM = "urn:zitadel:iam:org:project:376787351902355727:roles";

async function setup() {
  const { publicKey, privateKey } = await generateKeyPair("RS256");
  const verify = makeVerifier({
    issuer: ISS,
    audiences: ["aud-web", "proj-1"],
    requiredRole: "user",
    keyResolver: async () => publicKey,
  });
  const sign = (claims: Record<string, unknown>) =>
    new SignJWT(claims)
      .setProtectedHeader({ alg: "RS256" })
      .setIssuer(ISS)
      .setIssuedAt()
      .setExpirationTime("5m")
      .sign(privateKey);
  return { verify, sign };
}

describe("token verification", () => {
  it("accepts a token with the user role and a valid audience", async () => {
    const { verify, sign } = await setup();
    const t = await sign({ sub: "u1", aud: ["aud-web"], [ROLES_CLAIM]: { user: { org: "kuns.dev" } } });
    await expect(verify(t)).resolves.toMatchObject({ sub: "u1" });
  });

  it("accepts the role from the project-scoped claim variant", async () => {
    const { verify, sign } = await setup();
    const t = await sign({ sub: "u1", aud: "proj-1", [SCOPED_ROLES_CLAIM]: { user: { org: "kuns.dev" } } });
    await expect(verify(t)).resolves.toMatchObject({ sub: "u1" });
  });

  it("rejects a token without the required role", async () => {
    const { verify, sign } = await setup();
    const t = await sign({ sub: "u1", aud: ["aud-web"] });
    await expect(verify(t)).rejects.toThrow(/required role/);
  });

  it("rejects a token whose roles do not include the required one", async () => {
    const { verify, sign } = await setup();
    const t = await sign({ sub: "u1", aud: ["aud-web"], [ROLES_CLAIM]: { viewer: { org: "kuns.dev" } } });
    await expect(verify(t)).rejects.toThrow(/required role/);
  });

  it("rejects a token with no accepted audience", async () => {
    const { verify, sign } = await setup();
    const t = await sign({ sub: "u1", aud: ["other"], [ROLES_CLAIM]: { user: { org: "kuns.dev" } } });
    await expect(verify(t)).rejects.toThrow();
  });

  it("rejects a wrong issuer", async () => {
    const { publicKey, privateKey } = await generateKeyPair("RS256");
    const verify = makeVerifier({
      issuer: ISS,
      audiences: ["aud-web"],
      requiredRole: "user",
      keyResolver: async () => publicKey,
    });
    const t = await new SignJWT({ sub: "u1", aud: ["aud-web"], [ROLES_CLAIM]: { user: {} } })
      .setProtectedHeader({ alg: "RS256" })
      .setIssuer("https://evil.example")
      .setExpirationTime("5m")
      .sign(privateKey);
    await expect(verify(t)).rejects.toThrow();
  });
});

describe("rolesFromClaims", () => {
  it("merges generic and project-scoped role claims", () => {
    const roles = rolesFromClaims({
      [ROLES_CLAIM]: { user: {} },
      [SCOPED_ROLES_CLAIM]: { admin: {} },
    });
    expect(roles).toEqual(new Set(["user", "admin"]));
  });

  it("ignores malformed role claims and unrelated keys", () => {
    const roles = rolesFromClaims({
      [ROLES_CLAIM]: "not-an-object",
      "urn:zitadel:iam:org:project:roles:extra": { nope: {} },
      sub: "u1",
    } as never);
    expect(roles.size).toBe(0);
  });
});