import { NitroModules } from 'react-native-nitro-modules'
import { type Rsa, RsaEncryptionAlgorithmSpec } from './specs/Rsa.nitro'
import { bufferToInt8Array, int8ArrayToBuffer } from './utils/HybridInterop'

export enum RsaEncryptionAlgorithm {
  OaepWithSha1 = 'OaepWithSha1',
  OaepWithSha256 = 'OaepWithSha256',
}

export enum RsaSignatureAlgorithm {
  PssWithSha256 = 'PssWithSha256',
}

export type RsaAlgorithm = RsaEncryptionAlgorithm | RsaSignatureAlgorithm

function validateAlgorithm(alg: RsaAlgorithm) {
  switch (alg) {
    case RsaEncryptionAlgorithm.OaepWithSha1:
    case RsaEncryptionAlgorithm.OaepWithSha256:
    case RsaSignatureAlgorithm.PssWithSha256:
      break
    default:
      throw new Error(`Unsupported/invalid RSA algorithm ${alg}`)
  }
}

function validateAndGetHybridEncryptionAlgorithmSpec(
  key: XPrivateRsaKey | XPublicRsaKey
): RsaEncryptionAlgorithmSpec {
  switch (key.algorithm) {
    case RsaEncryptionAlgorithm.OaepWithSha1:
      return RsaEncryptionAlgorithmSpec.OaepWithSha1
    case RsaEncryptionAlgorithm.OaepWithSha256:
      return RsaEncryptionAlgorithmSpec.OaepWithSha256
    case RsaSignatureAlgorithm.PssWithSha256:
      throw new Error(
        `Invalid RSA key for encryption/decryption (key algorithm ${key.algorithm} is for signature/verification)`
      )
    default:
      throw new Error(
        `Unsupported/invalid RSA algorithm ${key.algorithm} in key`
      )
  }
}

function validateSignatureAlgorithm(key: XPrivateRsaKey | XPublicRsaKey) {
  switch (key.algorithm) {
    case RsaEncryptionAlgorithm.OaepWithSha1:
    case RsaEncryptionAlgorithm.OaepWithSha256:
      throw new Error(
        `Invalid RSA key for signature/verification (key algorithm ${key.algorithm} is for encryption/decryption)`
      )
    case RsaSignatureAlgorithm.PssWithSha256:
      break
    default:
      throw new Error(
        `Unsupported/invalid RSA algorithm ${key.algorithm} in key`
      )
  }
}

export enum RsaKeySize {
  Rsa2048 = 2048,
  Rsa4096 = 4096,
}

export interface XPrivateRsaKey {
  privateKey: Int8Array
  algorithm: RsaAlgorithm
}
export interface XPublicRsaKey {
  publicKey: Int8Array
  algorithm: RsaAlgorithm
}

export interface XRsaKeypair {
  readonly private: XPrivateRsaKey
  readonly public: XPublicRsaKey
}

export interface PartialXRsaService {
  generateKeyPair(
    algorithm: RsaAlgorithm,
    keySize?: RsaKeySize
  ): Promise<XRsaKeypair>
  exportPrivateKeyPkcs8(key: XPrivateRsaKey): Promise<Int8Array>
  exportPublicKeySpki(key: XPublicRsaKey): Promise<Int8Array>
  loadPrivateKeyPkcs8(
    algorithm: RsaAlgorithm,
    privateKeyPkcs8: Int8Array
  ): Promise<XPrivateRsaKey>
  loadPublicKeySpki(
    algorithm: RsaAlgorithm,
    publicKeySpki: Int8Array
  ): Promise<XPublicRsaKey>
  encrypt(data: Int8Array, publicKey: XPublicRsaKey): Promise<Int8Array>
  decrypt(data: Int8Array, privateKey: XPrivateRsaKey): Promise<Int8Array>
  sign(data: Int8Array, privateKey: XPrivateRsaKey): Promise<Int8Array>
  verifySignature(
    signature: Int8Array,
    data: Int8Array,
    publicKey: XPublicRsaKey
  ): Promise<boolean>
}

const hybridRsa = NitroModules.createHybridObject<Rsa>('Rsa')

const PEM_PRIVATE_HEADER = '-----BEGIN PRIVATE KEY-----\n'
const PEM_PRIVATE_FOOTER = '-----END PRIVATE KEY-----\n'
const PEM_PUBLIC_HEADER = '-----BEGIN PUBLIC KEY-----\n'
const PEM_PUBLIC_FOOTER = '-----END PUBLIC KEY-----\n'
const textEncoder = new TextEncoder()
const textDecoder = new TextDecoder()

function base64Encode(bytes: Int8Array): string {
  const asUbytes = new Uint8Array(bytes.buffer)
  let binary = ''
  for (let i = 0; i < asUbytes.length; i++) {
    binary += String.fromCharCode(asUbytes[i]!)
  }
  return btoa(binary)
}

function toPemString(
  bytes: Int8Array,
  header: string,
  footer: string
): Int8Array {
  const base64 = base64Encode(bytes)
  const lines = base64.match(/.{1,64}/g) ?? []
  const pem = header + lines.join('\n') + '\n' + footer
  return new Int8Array(textEncoder.encode(pem))
}

function unwrapPem(key: Int8Array, header: string, footer: string): Int8Array {
  const keyString = textDecoder.decode(key)

  if (!keyString.startsWith(header)) {
    throw new Error('Key does not start with expected pem header')
  }
  if (!keyString.endsWith(footer)) {
    throw new Error('Key does not end with expected pem footer')
  }

  const base64 = keyString
    .replace(header, '')
    .replace(footer, '')
    .replace(/\n/g, '')

  const binary = atob(base64)
  const bytes = new Int8Array(binary.length)
  const view = new Uint8Array(bytes.buffer)
  for (let i = 0; i < binary.length; i++) {
    view[i] = binary.charCodeAt(i)
  }

  return bytes
}

function wrapPkcs8ToPem(pkcs8: Int8Array): Int8Array {
  return toPemString(pkcs8, PEM_PRIVATE_HEADER, PEM_PRIVATE_FOOTER)
}

function unwrapPkcs8Pem(pem: Int8Array): Int8Array {
  return unwrapPem(pem, PEM_PRIVATE_HEADER, PEM_PRIVATE_FOOTER)
}

function wrapSpkiToPem(spki: Int8Array): Int8Array {
  return toPemString(spki, PEM_PUBLIC_HEADER, PEM_PUBLIC_FOOTER)
}

function unwrapSpkiPem(pem: Int8Array): Int8Array {
  return unwrapPem(pem, PEM_PUBLIC_HEADER, PEM_PUBLIC_FOOTER)
}

export const RsaService: PartialXRsaService = {
  async generateKeyPair(
    algorithm: RsaAlgorithm,
    keySize?: RsaKeySize
  ): Promise<XRsaKeypair> {
    validateAlgorithm(algorithm)
    if (
      keySize != undefined &&
      keySize !== RsaKeySize.Rsa4096 &&
      keySize !== RsaKeySize.Rsa2048
    )
      throw new Error(`Invalid/unsupported RSA key size ${keySize}`)
    const keypair = await hybridRsa.generateKeypair(
      keySize ?? RsaKeySize.Rsa2048
    )
    return {
      private: {
        privateKey: bufferToInt8Array(keypair.privateKey),
        algorithm: algorithm,
      },
      public: {
        publicKey: bufferToInt8Array(keypair.publicKey),
        algorithm: algorithm,
      },
    }
  },
  exportPrivateKeyPkcs8(key: XPrivateRsaKey): Promise<Int8Array> {
    return Promise.resolve(unwrapPkcs8Pem(key.privateKey))
  },
  exportPublicKeySpki(key: XPublicRsaKey): Promise<Int8Array> {
    return Promise.resolve(unwrapSpkiPem(key.publicKey))
  },
  loadPrivateKeyPkcs8(
    algorithm: RsaAlgorithm,
    privateKeyPkcs8: Int8Array
  ): Promise<XPrivateRsaKey> {
    validateAlgorithm(algorithm)
    const pem = wrapPkcs8ToPem(privateKeyPkcs8)
    try {
      hybridRsa.checkValidPrivate(int8ArrayToBuffer(pem))
    } catch (e) {
      throw new Error('Invalid pkcs8 key', { cause: e })
    }
    return Promise.resolve({
      algorithm,
      privateKey: pem,
    })
  },
  loadPublicKeySpki(
    algorithm: RsaAlgorithm,
    publicKeySpki: Int8Array
  ): Promise<XPublicRsaKey> {
    validateAlgorithm(algorithm)
    const pem = wrapSpkiToPem(publicKeySpki)
    try {
      hybridRsa.checkValidPublic(int8ArrayToBuffer(pem))
    } catch (e) {
      throw new Error('Invalid spki key', { cause: e })
    }
    return Promise.resolve({
      algorithm,
      publicKey: pem,
    })
  },
  async sign(data: Int8Array, privateKey: XPrivateRsaKey): Promise<Int8Array> {
    validateSignatureAlgorithm(privateKey)
    const res = await hybridRsa.sign(
      int8ArrayToBuffer(data),
      int8ArrayToBuffer(privateKey.privateKey)
    )
    return bufferToInt8Array(res)
  },
  async verifySignature(
    signature: Int8Array,
    data: Int8Array,
    publicKey: XPublicRsaKey
  ): Promise<boolean> {
    validateSignatureAlgorithm(publicKey)
    return await hybridRsa.verify(
      int8ArrayToBuffer(signature),
      int8ArrayToBuffer(data),
      int8ArrayToBuffer(publicKey.publicKey)
    )
  },
  async decrypt(
    data: Int8Array,
    privateKey: XPrivateRsaKey
  ): Promise<Int8Array> {
    const res = await hybridRsa.decrypt(
      validateAndGetHybridEncryptionAlgorithmSpec(privateKey),
      int8ArrayToBuffer(data),
      int8ArrayToBuffer(privateKey.privateKey)
    )
    return bufferToInt8Array(res)
  },
  async encrypt(data: Int8Array, publicKey: XPublicRsaKey): Promise<Int8Array> {
    const res = await hybridRsa.encrypt(
      validateAndGetHybridEncryptionAlgorithmSpec(publicKey),
      int8ArrayToBuffer(data),
      int8ArrayToBuffer(publicKey.publicKey)
    )
    return bufferToInt8Array(res)
  },
}