// Helper function to Base64Url encode a string
function base64UrlEncode(str: string): string {
  return btoa(str).replace(/\+/g, "-").replace(/\//g, "_").replace(/=+$/, "");
}

// Helper function to Base64Url decode a string
function base64UrlDecode(str: string): string {
  // Replace URL-safe characters with Base64 standard characters
  str = str.replace(/-/g, "+").replace(/_/g, "/");
  // Add padding characters if necessary
  switch (str.length % 4) {
    case 2:
      str += "==";
      break;
    case 3:
      str += "=";
      break;
  }
  return atob(str);
}

type JwtHeader = {
  alg: string;
  typ: string;
};

const HEADER: JwtHeader = {
  alg: "HS256",
  typ: "JWT",
};

const JWT_ERROR = {
  INVALID_JWT_FORMAT: "Invalid JWT format",
  // INVALID_SIGNATURE: "Invalid signature",
};

// Function to create a JWT
async function createJwt<
  P extends Record<string, unknown>,
  H extends JwtHeader = JwtHeader,
>({
  payload,
  secret,
  header,
}: {
  payload: P;
  secret: string;
  header?: H;
}): Promise<string> {
  const encoder = new TextEncoder();

  // Step 1: Encode Header
  const encodedHeader = base64UrlEncode(
    JSON.stringify({ ...HEADER, ...header })
  );

  // Step 2: Encode Payload
  const encodedPayload = base64UrlEncode(JSON.stringify(payload));

  // Step 3: Create Data to Sign
  const dataToSign = `${encodedHeader}.${encodedPayload}`;

  // Step 4: Convert Secret to CryptoKey
  const key = await crypto.subtle.importKey(
    "raw",
    encoder.encode(secret),
    { name: "HMAC", hash: "SHA-256" },
    false,
    ["sign"]
  );

  // Step 5: Sign Data
  const signature = await crypto.subtle.sign(
    "HMAC",
    key,
    encoder.encode(dataToSign)
  );

  // Step 6: Base64Url Encode Signature
  const encodedSignature = base64UrlEncode(
    String.fromCharCode(...new Uint8Array(signature))
  );

  // Step 7: Assemble JWT
  return `${dataToSign}.${encodedSignature}`;
}

// Function to decode and verify a JWT
async function decodeJwt<
  P extends Record<string, unknown>,
  H extends JwtHeader = JwtHeader,
>({
  token,
  secret,
}: {
  token: string;
  secret?: string;
}): Promise<{ header: H; payload: P; valid: boolean }> {
  if (typeof token !== "string") {
    throw new Error(JWT_ERROR.INVALID_JWT_FORMAT);
  }

  // Split the JWT into its parts
  const [encodedHeader, encodedPayload, encodedSignature] = token.split(".");

  if (!encodedHeader || !encodedPayload || !encodedSignature) {
    throw new Error(JWT_ERROR.INVALID_JWT_FORMAT);
  }

  // Base64Url decode the header and payload
  const headerJson = base64UrlDecode(encodedHeader);
  const payloadJson = base64UrlDecode(encodedPayload);

  const header: H = JSON.parse(headerJson);
  const payload: P = JSON.parse(payloadJson);

  // Recreate the signature data
  const dataToVerify = `${encodedHeader}.${encodedPayload}`;

  const encoder = new TextEncoder();

  if (secret !== undefined) {
    // Convert the secret to CryptoKey
    const key = await crypto.subtle.importKey(
      "raw",
      encoder.encode(secret),
      { name: "HMAC", hash: "SHA-256" },
      false,
      ["verify"]
    );

    // Convert the Base64Url encoded signature to a Uint8Array
    const binarySignature = base64UrlDecode(encodedSignature);
    const signature = new Uint8Array(
      [...binarySignature].map((char) => char.charCodeAt(0))
    );

    // Verify the signature
    const valid = await crypto.subtle.verify(
      "HMAC",
      key,
      signature,
      encoder.encode(dataToVerify)
    );

    // if (!valid) {
    //   throw new Error(JWT_ERROR.INVALID_SIGNATURE);
    // }

    return { header, payload, valid };
  }

  return { header, payload, valid: false };
}

export namespace jwt {
  export const encode = createJwt;
  export const decode = decodeJwt;
  export const ERROR = JWT_ERROR;
  export const DEFAULT_HEADER = HEADER;
}