const encoder = new TextEncoder();
const RSA_BITS = 4096;
const PRIME_BITS = 2048;
const PUBLIC_EXPONENT = 65537n;

export const DERIVATION_NAME = 'SAWD-v1';
export const MASTER_LABEL = 'SHINE/ARWEAVE/RSA4096/SAWD-v1/MASTER';
export const STREAM_LABEL = 'SHINE/ARWEAVE/RSA4096/SAWD-v1/STREAM';
export const MR_LABEL = 'SHINE/ARWEAVE/RSA4096/SAWD-v1/MILLER-RABIN';
export const MILLER_RABIN_ROUNDS = 42;
export const SMALL_PRIME_LIMIT = 10000;

function getSubtle() {
  const subtle = globalThis.crypto?.subtle;
  if (!subtle) {
    throw new Error('SAWD-v1 требует WebCrypto (crypto.subtle)');
  }
  return subtle;
}

function utf8(text) {
  return encoder.encode(String(text));
}

function concatBytes(...chunks) {
  const total = chunks.reduce((sum, part) => sum + (part?.length || 0), 0);
  const out = new Uint8Array(total);
  let offset = 0;
  for (const part of chunks) {
    if (!part?.length) continue;
    out.set(part, offset);
    offset += part.length;
  }
  return out;
}

function bytesToBase64(bytes) {
  let binary = '';
  const chunkSize = 0x8000;
  for (let i = 0; i < bytes.length; i += chunkSize) {
    const chunk = bytes.subarray(i, i + chunkSize);
    binary += String.fromCharCode(...chunk);
  }
  return btoa(binary);
}

function base64UrlToBytes(value) {
  const normalized = String(value || '').replace(/-/g, '+').replace(/_/g, '/');
  const padLen = (4 - (normalized.length % 4)) % 4;
  const binary = atob(normalized + '='.repeat(padLen));
  const out = new Uint8Array(binary.length);
  for (let i = 0; i < binary.length; i += 1) {
    out[i] = binary.charCodeAt(i);
  }
  return out;
}

export function base64UrlEncode(bytes) {
  return bytesToBase64(bytes).replace(/\+/g, '-').replace(/\//g, '_').replace(/=+$/g, '');
}

export function unsignedBytesToBigInt(bytes) {
  let value = 0n;
  for (const byte of bytes) {
    value = (value << 8n) | BigInt(byte);
  }
  return value;
}

export function bigIntToUnsignedBytes(value) {
  let v = BigInt(value);
  if (v < 0n) throw new Error('Ожидалось неотрицательное BigInt значение');
  if (v === 0n) return new Uint8Array([0]);

  let hex = v.toString(16);
  if (hex.length % 2 !== 0) hex = `0${hex}`;
  const out = new Uint8Array(hex.length / 2);
  for (let i = 0; i < out.length; i += 1) {
    out[i] = Number.parseInt(hex.slice(i * 2, i * 2 + 2), 16);
  }
  return out;
}

function bitLength(v) {
  return BigInt(v).toString(2).length;
}

function uint64be(counter) {
  const out = new Uint8Array(8);
  let v = BigInt(counter);
  for (let i = 7; i >= 0; i -= 1) {
    out[i] = Number(v & 0xffn);
    v >>= 8n;
  }
  return out;
}

function uint32be(counter) {
  const out = new Uint8Array(4);
  let v = BigInt(counter);
  for (let i = 3; i >= 0; i -= 1) {
    out[i] = Number(v & 0xffn);
    v >>= 8n;
  }
  return out;
}

async function importHmacKey(keyBytes) {
  return getSubtle().importKey(
    'raw',
    keyBytes,
    { name: 'HMAC', hash: 'SHA-256' },
    false,
    ['sign'],
  );
}

async function hmacSha256WithImportedKey(importedKey, messageBytes) {
  const signed = await getSubtle().sign('HMAC', importedKey, messageBytes);
  return new Uint8Array(signed);
}

export async function hmacSha256(keyBytes, messageBytes) {
  const imported = await importHmacKey(keyBytes);
  return hmacSha256WithImportedKey(imported, messageBytes);
}

export async function sha256(bytes) {
  const digest = await getSubtle().digest('SHA-256', bytes);
  return new Uint8Array(digest);
}

const STREAM_PREFIX = utf8(`${STREAM_LABEL}/`);
const MR_PREFIX = utf8(`${MR_LABEL}/`);
const SLASH = utf8('/');
const MASTER_LABEL_UTF8 = utf8(MASTER_LABEL);
const TWO = 2n;
const THREE = 3n;
const PRIME_HIGH_BIT_MASK = 1n << 2047n;

function makeSmallPrimes(limit) {
  const sieve = new Uint8Array(limit + 1);
  const primes = [];
  for (let i = 2; i <= limit; i += 1) {
    if (sieve[i]) continue;
    primes.push(i);
    for (let j = i * 2; j <= limit; j += i) {
      sieve[j] = 1;
    }
  }
  return primes;
}

const SMALL_PRIMES = makeSmallPrimes(SMALL_PRIME_LIMIT);

export async function deriveBytes(masterSeed32, label, length) {
  const imported = await importHmacKey(masterSeed32);
  return deriveBytesWithImportedKey(imported, label, length);
}

async function deriveBytesWithImportedKey(masterSeedKey, label, length) {
  const target = Number(length);
  if (!Number.isInteger(target) || target <= 0) {
    throw new Error('deriveBytes: length должен быть положительным целым');
  }
  let counter = 0n;
  const chunks = [];
  let written = 0;
  const labelBytes = utf8(String(label || ''));
  while (written < target) {
    const msg = concatBytes(STREAM_PREFIX, labelBytes, SLASH, uint64be(counter));
    const block = await hmacSha256WithImportedKey(masterSeedKey, msg);
    chunks.push(block);
    written += block.length;
    counter += 1n;
  }
  return concatBytes(...chunks).slice(0, target);
}

export function gcd(a, b) {
  let x = BigInt(a);
  let y = BigInt(b);
  while (y !== 0n) {
    const t = x % y;
    x = y;
    y = t;
  }
  return x < 0n ? -x : x;
}

export function lcm(a, b) {
  const x = BigInt(a);
  const y = BigInt(b);
  return (x / gcd(x, y)) * y;
}

export function modPow(base, exponent, modulus) {
  let b = BigInt(base) % BigInt(modulus);
  let e = BigInt(exponent);
  const m = BigInt(modulus);
  let result = 1n;
  while (e > 0n) {
    if (e & 1n) result = (result * b) % m;
    b = (b * b) % m;
    e >>= 1n;
  }
  return result;
}

export function modInverse(a, modulus) {
  let t = 0n;
  let newT = 1n;
  let r = BigInt(modulus);
  let newR = BigInt(a) % BigInt(modulus);

  while (newR !== 0n) {
    const q = r / newR;
    [t, newT] = [newT, t - q * newT];
    [r, newR] = [newR, r - q * newR];
  }
  if (r !== 1n) throw new Error('Обратный элемент не существует');
  if (t < 0n) t += BigInt(modulus);
  return t;
}

function passesSmallPrimeFilter(candidate) {
  for (const p of SMALL_PRIMES) {
    const pBig = BigInt(p);
    if (candidate === pBig) return true;
    if (candidate % pBig === 0n) return false;
  }
  return true;
}

async function millerRabinWithImportedKey(masterSeedKey, candidate, label, index, rounds = MILLER_RABIN_ROUNDS) {
  const n = BigInt(candidate);
  if (n < 2n) return false;
  if (n === 2n || n === 3n) return true;
  if ((n & 1n) === 0n) return false;

  let d = n - 1n;
  let s = 0;
  while ((d & 1n) === 0n) {
    d >>= 1n;
    s += 1;
  }

  const labelBytes = utf8(String(label || ''));
  const indexBytes = uint64be(index);
  const nMinusThree = n - THREE;
  if (nMinusThree <= 0n) return false;

  for (let round = 0; round < rounds; round += 1) {
    const msg = concatBytes(MR_PREFIX, labelBytes, SLASH, indexBytes, SLASH, uint32be(round));
    const baseBytes = await hmacSha256WithImportedKey(masterSeedKey, msg);
    const a = TWO + (unsignedBytesToBigInt(baseBytes) % nMinusThree);

    let x = modPow(a, d, n);
    if (x === 1n || x === n - 1n) continue;

    let witnessComposite = true;
    for (let i = 1; i < s; i += 1) {
      x = modPow(x, 2n, n);
      if (x === n - 1n) {
        witnessComposite = false;
        break;
      }
    }
    if (witnessComposite) return false;
  }

  return true;
}

export async function millerRabin(masterSeed32, candidate, label, index, rounds = MILLER_RABIN_ROUNDS) {
  const imported = await importHmacKey(masterSeed32);
  return millerRabinWithImportedKey(imported, candidate, label, index, rounds);
}

export async function isProbablePrime(masterSeed32, candidate, label, index) {
  const imported = await importHmacKey(masterSeed32);
  return isProbablePrimeWithImportedKey(imported, candidate, label, index);
}

async function isProbablePrimeWithImportedKey(masterSeedKey, candidate, label, index) {
  const n = BigInt(candidate);
  if ((n & 1n) === 0n) return false;
  if (bitLength(n) !== PRIME_BITS) return false;
  if (!passesSmallPrimeFilter(n)) return false;
  if (gcd(n - 1n, PUBLIC_EXPONENT) !== 1n) return false;
  return millerRabinWithImportedKey(masterSeedKey, n, label, index, MILLER_RABIN_ROUNDS);
}

async function tickEvery(iteration, step = 8n) {
  if (iteration % step === 0n) {
    await Promise.resolve();
  }
}

export async function derivePrime(masterSeed32, label) {
  const imported = await importHmacKey(masterSeed32);
  return derivePrimeWithImportedKey(imported, label);
}

async function derivePrimeWithImportedKey(masterSeedKey, label, startIndex = 0n) {
  let index = BigInt(startIndex);
  while (true) {
    const raw = await deriveBytesWithImportedKey(masterSeedKey, `${label}/${index}`, PRIME_BITS / 8);
    let candidate = unsignedBytesToBigInt(raw);
    candidate |= PRIME_HIGH_BIT_MASK;
    candidate |= 1n;

    if (bitLength(candidate) === PRIME_BITS && await isProbablePrimeWithImportedKey(masterSeedKey, candidate, label, index)) {
      return { prime: candidate, index };
    }

    index += 1n;
    await tickEvery(index);
  }
}

function toJwkB64(value) {
  return base64UrlEncode(bigIntToUnsignedBytes(value));
}

async function deriveArweaveWalletParts(deviceKey32) {
  if (!(deviceKey32 instanceof Uint8Array)) {
    throw new Error('SAWD-v1: deviceKey32 должен быть Uint8Array');
  }
  if (deviceKey32.length !== 32) {
    throw new Error('SAWD-v1: deviceKey32 должен быть ровно 32 байта');
  }

  const masterSeed32 = await hmacSha256(MASTER_LABEL_UTF8, deviceKey32);
  const masterSeedKey = await importHmacKey(masterSeed32);

  const pResult = await derivePrimeWithImportedKey(masterSeedKey, 'p');
  let qResult = await derivePrimeWithImportedKey(masterSeedKey, 'q');
  while (qResult.prime === pResult.prime) {
    qResult = await derivePrimeWithImportedKey(masterSeedKey, 'q', qResult.index + 1n);
  }

  let p = pResult.prime;
  let q = qResult.prime;
  if (p > q) {
    [p, q] = [q, p];
  }

  const n = p * q;
  const lambda = lcm(p - 1n, q - 1n);
  const d = modInverse(PUBLIC_EXPONENT, lambda);
  const dp = d % (p - 1n);
  const dq = d % (q - 1n);
  const qi = modInverse(q, p);

  const jwk = {
    kty: 'RSA',
    e: 'AQAB',
    n: toJwkB64(n),
    d: toJwkB64(d),
    p: toJwkB64(p),
    q: toJwkB64(q),
    dp: toJwkB64(dp),
    dq: toJwkB64(dq),
    qi: toJwkB64(qi),
  };

  const owner = jwk.n;
  const address = base64UrlEncode(await sha256(bigIntToUnsignedBytes(n)));

  return {
    derivation: DERIVATION_NAME,
    jwk,
    owner,
    address,
    _private: {
      p,
      q,
      n,
    },
  };
}

export async function deriveArweaveWalletFromDeviceKey32(deviceKey32) {
  const result = await deriveArweaveWalletParts(deviceKey32);
  return {
    derivation: result.derivation,
    jwk: result.jwk,
    owner: result.owner,
    address: result.address,
  };
}

export async function selfTestSawdV1() {
  const invalid = new Uint8Array(31);
  let invalidFailed = false;
  try {
    await deriveArweaveWalletFromDeviceKey32(invalid);
  } catch {
    invalidFailed = true;
  }
  if (!invalidFailed) {
    throw new Error('SAWD-v1 self-test: длина != 32 должна приводить к ошибке');
  }

  const deviceKey = new Uint8Array(32);
  for (let i = 0; i < deviceKey.length; i += 1) {
    deviceKey[i] = i + 1;
  }

  const first = await deriveArweaveWalletParts(deviceKey);
  const second = await deriveArweaveWalletParts(deviceKey);

  if (!first.address || !second.address) {
    throw new Error('SAWD-v1 self-test: адрес пустой');
  }
  if (first.address !== second.address) {
    throw new Error('SAWD-v1 self-test: адрес должен быть детерминированным');
  }
  if (first.jwk.n !== second.jwk.n) {
    throw new Error('SAWD-v1 self-test: jwk.n должен быть детерминированным');
  }
  if (first.jwk.e !== 'AQAB') {
    throw new Error('SAWD-v1 self-test: jwk.e должен быть AQAB');
  }
  if (first.owner !== first.jwk.n) {
    throw new Error('SAWD-v1 self-test: owner должен совпадать с jwk.n');
  }
  if (!(first._private.p < first._private.q)) {
    throw new Error('SAWD-v1 self-test: ожидается p < q');
  }
  if (bitLength(first._private.n) !== RSA_BITS) {
    throw new Error('SAWD-v1 self-test: ожидается n с длиной 4096 бит');
  }

  const pFromJwk = unsignedBytesToBigInt(base64UrlToBytes(first.jwk.p));
  const qFromJwk = unsignedBytesToBigInt(base64UrlToBytes(first.jwk.q));
  if (!(pFromJwk < qFromJwk)) {
    throw new Error('SAWD-v1 self-test: p и q в JWK должны удовлетворять p < q');
  }

  return {
    ok: true,
    derivation: DERIVATION_NAME,
    address: first.address,
  };
}
