import {
  JwtError,
  signJwt,
  verifyJwt,
  type JwtClaims,
  type SignOptions,
  type VerifyOptions,
} from "./index.ts";

export interface JwtKey {
  id: string;
  secret: string;
  active?: boolean;
}

export interface JwtKeyring {
  active(): JwtKey;
  resolve(id: string): JwtKey | undefined;
  keys(): JwtKey[];
}

const KEY_ID = /^[A-Za-z0-9._-]{1,64}$/;

function decodePart(value: string): Record<string, unknown> {
  const padding = value.length % 4 === 0 ? "" : "=".repeat(4 - (value.length % 4));
  const json = atob(value.replace(/-/g, "+").replace(/_/g, "/") + padding);
  return JSON.parse(
    new TextDecoder().decode(Uint8Array.from(json, (character) => character.charCodeAt(0))),
  ) as Record<string, unknown>;
}

export function decodeJwt(token: string): {
  header: Record<string, unknown>;
  claims: JwtClaims;
} {
  const parts = token.split(".");
  if (parts.length !== 3) throw new JwtError("Malformed token");
  try {
    return {
      header: decodePart(parts[0]!),
      claims: decodePart(parts[1]!) as JwtClaims,
    };
  } catch {
    throw new JwtError("Invalid token encoding");
  }
}

export function createJwtKeyring(keys: JwtKey[]): JwtKeyring {
  const values = new Map<string, JwtKey>();
  let activeCount = 0;
  for (const key of keys) {
    if (!KEY_ID.test(key.id)) throw new TypeError(`Invalid JWT key id: ${key.id}`);
    if (!key.secret.trim()) throw new TypeError(`JWT key '${key.id}' has an empty secret`);
    if (values.has(key.id)) throw new Error(`WRN-JWT-KEYRING-DUPLICATE: ${key.id}`);
    if (key.active) activeCount++;
    values.set(key.id, { ...key });
  }
  if (!values.size) throw new Error("WRN-JWT-KEYRING-EMPTY");
  if (activeCount > 1) throw new Error("WRN-JWT-KEYRING-MULTIPLE-ACTIVE");

  const activeInternal = (): JwtKey =>
    [...values.values()].find((key) => key.active) ?? [...values.values()].at(-1)!;

  return {
    active: () => ({ ...activeInternal() }),
    resolve(id) {
      const key = values.get(id);
      return key ? { ...key } : undefined;
    },
    keys: () => [...values.values()].map((key) => ({ ...key })),
  };
}

export async function signWithKeyring(
  claims: JwtClaims,
  keyring: JwtKeyring,
  options: SignOptions = {},
): Promise<string> {
  const key = keyring.active();
  return signJwt(claims, key.secret, { ...options, keyId: key.id });
}

export async function verifyWithKeyring<T extends JwtClaims = JwtClaims>(
  token: string,
  keyring: JwtKeyring,
  options: VerifyOptions = {},
): Promise<T> {
  const { header } = decodeJwt(token);
  const kid = typeof header.kid === "string" ? header.kid : undefined;
  const keys = keyring.keys();
  const key = kid ? keyring.resolve(kid) : keys.length === 1 ? keys[0] : undefined;
  if (!key) {
    throw new JwtError(kid ? `Unknown key id: ${kid}` : "Token has no key id");
  }
  return verifyJwt<T>(token, key.secret, options);
}