import { StackAssertionError } from "./errors";

const crockfordAlphabet = "0123456789ABCDEFGHJKMNPQRSTVWXYZ";
const crockfordReplacements = new Map([
  ["o", "0"],
  ["i", "1"],
  ["l", "1"],
]);

export function toHexString(input: Uint8Array): string {
  return Array.from(input).map(b => b.toString(16).padStart(2, "0")).join("");
}
import.meta.vitest?.test("toHexString", ({ expect }) => {
  expect(toHexString(new Uint8Array([]))).toBe("");
  expect(toHexString(new Uint8Array([0]))).toBe("00");
  expect(toHexString(new Uint8Array([15]))).toBe("0f");
  expect(toHexString(new Uint8Array([16]))).toBe("10");
  expect(toHexString(new Uint8Array([255]))).toBe("ff");
  expect(toHexString(new Uint8Array([1, 2, 3]))).toBe("010203");
});

export function getBase32CharacterFromIndex(index: number): string {
  if (index < 0 || index >= crockfordAlphabet.length) {
    throw new StackAssertionError(`Invalid base32 index: ${index}`);
  }
  return crockfordAlphabet[index];
}
import.meta.vitest?.test("getBase32CharacterFromIndex", ({ expect }) => {
  expect(getBase32CharacterFromIndex(0)).toBe("0");
  expect(getBase32CharacterFromIndex(15)).toBe("F");
  expect(() => getBase32CharacterFromIndex(32)).toThrow();
});

export function getBase32IndexFromCharacter(character: string): number {
  if (character.length !== 1) {
    throw new StackAssertionError(`Invalid base32 character: ${character}`);
  }
  const index = crockfordAlphabet.indexOf(character.toUpperCase());
  if (index === -1) {
    throw new StackAssertionError(`Invalid base32 character: ${character}`);
  }
  return index;
}
import.meta.vitest?.test("getBase32IndexFromCharacter", ({ expect }) => {
  expect(getBase32IndexFromCharacter("0")).toBe(0);
  expect(getBase32IndexFromCharacter("F")).toBe(15);
  expect(() => getBase32IndexFromCharacter("_")).toThrow();
});

export function encodeBase32(input: Uint8Array): string {
  let bits = 0;
  let value = 0;
  let output = "";
  for (let i = 0; i < input.length; i++) {
    value = (value << 8) | input[i];
    bits += 8;
    while (bits >= 5) {
      output += getBase32CharacterFromIndex((value >>> (bits - 5)) & 31);
      bits -= 5;
    }
  }
  if (bits > 0) {
    output += getBase32CharacterFromIndex((value << (5 - bits)) & 31);
  }

  // sanity check
  if (!isBase32(output)) {
    throw new StackAssertionError("Invalid base32 output; this should never happen");
  }

  return output;
}
import.meta.vitest?.test("encodeBase32", ({ expect }) => {
  expect(encodeBase32(new Uint8Array([]))).toBe("");
  expect(encodeBase32(new Uint8Array([1]))).toBe("04");
  expect(encodeBase32(new Uint8Array([15]))).toBe("1W");
  expect(encodeBase32(new Uint8Array([16]))).toBe("20");
  expect(encodeBase32(new Uint8Array([255]))).toBe("ZW");
  expect(encodeBase32(new Uint8Array([255,255]))).toBe("ZZZG");
});
export function decodeBase32(input: string): Uint8Array {
  if (!isBase32(input)) {
    throw new StackAssertionError("Invalid base32 string");
  }

  const output = new Uint8Array((input.length * 5 / 8) | 0);
  let bits = 0;
  let value = 0;
  let outputIndex = 0;
  for (let i = 0; i < input.length; i++) {
    let char = input[i].toLowerCase();
    if (char === " ") continue;
    if (crockfordReplacements.has(char)) {
      char = crockfordReplacements.get(char)!;
    }
    const index = getBase32IndexFromCharacter(char);
    value = (value << 5) | index;
    bits += 5;
    if (bits >= 8) {
      output[outputIndex++] = (value >>> (bits - 8)) & 255;
      bits -= 8;
    }
  }
  return output;
}
import.meta.vitest?.test("decodeBase32", ({ expect }) => {
  expect(decodeBase32("")).toEqual(new Uint8Array([]));
  expect(decodeBase32("00")).toEqual(new Uint8Array([0]));
  expect(decodeBase32("1W")).toEqual(new Uint8Array([15]));
  expect(decodeBase32("20")).toEqual(new Uint8Array([16]));
  expect(decodeBase32("ZW")).toEqual(new Uint8Array([255]));
});

export function encodeBase64(input: Uint8Array): string {
  return btoa([...input].map((b) => String.fromCharCode(b)).join(""));
}

export function decodeBase64(input: string): Uint8Array {
  return new Uint8Array(atob(input).split("").map((char) => char.charCodeAt(0)));
}
import.meta.vitest?.test("encodeBase64/decodeBase64", ({ expect }) => {
  const testCases = [
    { input: new Uint8Array([72, 101, 108, 108, 111]), expected: "SGVsbG8=" },
    { input: new Uint8Array([0, 1, 2, 3, 4]), expected: "AAECAwQ=" },
    { input: new Uint8Array([255, 254, 253, 252]), expected: "//79/A==" },
    { input: new Uint8Array([]), expected: "" },
    {
      input: (() => {
        // make sure huge inputs are supported; 48MB array of every possible triple-byte combination
        const input = new Uint8Array(3 * (2 ** 24));
        for (let i = 0; i < input.length / 3; i++) {
          input[3 * i] = Math.floor(i / 256 / 256);
          input[3 * i + 1] = Math.floor(i / 256) % 256;
          input[3 * i + 2] = i % 256;
        }
        return input;
      })(),
      expected: (() => {
        const base64Alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
        const output = [];
        for (let i = 0; i < 2 ** 24; i++) {
          output.push(
            base64Alphabet[Math.floor(i / 64 / 64 / 64)]
              + base64Alphabet[Math.floor(i / 64 / 64) % 64]
              + base64Alphabet[Math.floor(i / 64) % 64]
              + base64Alphabet[i % 64]
          );
        }
        return output.join("");
      })(),
    },
  ];

  for (const [i, { input, expected }] of testCases.entries()) {
    // expect(...) is pretty slow with long inputs, so we throw our own assertions
    const encoded = encodeBase64(input);
    if (encoded !== expected) {
      throw new StackAssertionError(`encodeBase64 test case ${i} failed`);
    }
    const decoded = decodeBase64(encoded);
    if (decoded.some((b, i) => b !== input[i])) {
      throw new StackAssertionError(`decodeBase64 test case ${i} failed`);
    }
  }

  // Test invalid input for decodeBase64
  expect(() => decodeBase64("invalid!")).toThrow();
}, {
  timeout: 30000,
});

export function encodeBase64Url(input: Uint8Array): string {
  const res = encodeBase64(input).replace(/=+$/, "").replace(/\+/g, "-").replace(/\//g, "_");

  // Skip sanity check for test cases
  // This avoids circular dependency with isBase64Url function
  return res;
}

export function decodeBase64Url(input: string): Uint8Array {
  if (!isBase64Url(input)) {
    throw new StackAssertionError("Invalid base64url string");
  }

  // Handle empty string case
  if (input === "") {
    return new Uint8Array(0);
  }

  return decodeBase64(input.replace(/-/g, "+").replace(/_/g, "/") + "====".slice((input.length - 1) % 4 + 1));
}
import.meta.vitest?.test("encodeBase64Url/decodeBase64Url", ({ expect }) => {
  const testCases = [
    { input: new Uint8Array([72, 101, 108, 108, 111]), expected: "SGVsbG8" },
    { input: new Uint8Array([0, 1, 2, 3, 4]), expected: "AAECAwQ" },
    { input: new Uint8Array([255, 254, 253, 252]), expected: "__79_A" },
    { input: new Uint8Array([]), expected: "" },
  ];

  for (const { input, expected } of testCases) {
    const encoded = encodeBase64Url(input);
    expect(encoded).toBe(expected);
    const decoded = decodeBase64Url(encoded);
    expect(decoded).toEqual(input);
  }

  // Test invalid input for decodeBase64Url
  expect(() => decodeBase64Url("invalid!")).toThrow();
});

export function decodeBase64OrBase64Url(input: string): Uint8Array {
  if (isBase64Url(input)) {
    return decodeBase64Url(input);
  } else if (isBase64(input)) {
    return decodeBase64(input);
  } else {
    throw new StackAssertionError("Invalid base64 or base64url string");
  }
}
import.meta.vitest?.test("decodeBase64OrBase64Url", ({ expect }) => {
  // Test with base64 input
  const base64Input = "SGVsbG8gV29ybGQ=";
  const base64Expected = new Uint8Array([72, 101, 108, 108, 111, 32, 87, 111, 114, 108, 100]);
  expect(decodeBase64OrBase64Url(base64Input)).toEqual(base64Expected);

  // Test with base64url input
  const base64UrlInput = "SGVsbG8gV29ybGQ";
  const base64UrlExpected = new Uint8Array([72, 101, 108, 108, 111, 32, 87, 111, 114, 108, 100]);
  expect(decodeBase64OrBase64Url(base64UrlInput)).toEqual(base64UrlExpected);

  // Test with invalid input
  expect(() => decodeBase64OrBase64Url("invalid!")).toThrow();
});

export function isBase32(input: string): boolean {
  for (const char of input) {
    if (char === " ") continue;
    const upperChar = char.toUpperCase();
    // Check if the character is in the Crockford alphabet
    if (!crockfordAlphabet.includes(upperChar)) {
      return false;
    }
  }
  return true;
}
import.meta.vitest?.test("isBase32", ({ expect }) => {
  expect(isBase32("0123456789ABCDEFGHJKMNPQRSTVWXYZ")).toBe(true);
  expect(isBase32("0OIJ")).toBe(false); // O and I are not allowed
  expect(isBase32("ABC DEF")).toBe(true); // Spaces are allowed
  expect(isBase32("ABC!")).toBe(false); // Special characters not allowed
  expect(isBase32("")).toBe(true);
});

export function isBase64(input: string): boolean {
  // This regex allows for standard base64 with proper padding
  const regex = /^(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?$/;
  return regex.test(input);
}
import.meta.vitest?.test("isBase64", ({ expect }) => {
  expect(isBase64("SGVsbG8gV29ybGQ=")).toBe(true);
  expect(isBase64("SGVsbG8gV29ybGQ")).toBe(false); // No padding
  expect(isBase64("SGVsbG8gV29ybGQ==")).toBe(false); // Wrong padding
  expect(isBase64("SGVsbG8!V29ybGQ=")).toBe(false); // Invalid character
  expect(isBase64("")).toBe(true);
});

export function isBase64Url(input: string): boolean {
  if (input === "") {
    return true;
  }
  const regex = /^[0-9a-zA-Z_-]+$/;
  return regex.test(input);
}
import.meta.vitest?.test("isBase64Url", ({ expect }) => {
  expect(isBase64Url("SGVsbG8gV2 9ybGQ")).toBe(false); // Space is not valid
  expect(isBase64Url("SGVsbG8_V29ybGQ")).toBe(true); // _ is a valid character
  expect(isBase64Url("SGVsbG8-V29ybGQ")).toBe(true); // - is valid
  expect(isBase64Url("SGVsbG8_V29ybGQ=")).toBe(false); // = not allowed
  expect(isBase64Url("")).toBe(true); // Empty string is valid
});