import type * as types from '../types.d.ts'
import { invalidKeyInput } from './invalid_key_input.js'
import { encodeBase64, decodeBase64 } from '../lib/base64.js'
import { JOSENotSupported } from '../util/errors.js'
import { keyAlgorithm, unsupportedAlg, algArgument } from './key_algorithm.js'
import { isCryptoKey, isKeyObject } from './is_key_like.js'

import type { KeyImportOptions } from '../key/import.js'
import { validateExtractableOption } from './key_options.js'

/**
 * Formats a base64 string as a PEM-encoded key with proper line breaks and headers.
 *
 * @param b64 - Base64-encoded key data
 * @param descriptor - Key type descriptor (e.g., "PUBLIC KEY", "PRIVATE KEY")
 *
 * @returns PEM-formatted string
 */
const formatPEM = (b64: string, descriptor: string) => {
  const newlined = (b64.match(/.{1,64}/g) || []).join('\n')
  return `-----BEGIN ${descriptor}-----\n${newlined}\n-----END ${descriptor}-----`
}

interface ExportOptions {
  format: 'pem'
  type: 'spki' | 'pkcs8'
}

interface ExtractableKeyObject extends types.KeyObject {
  export(arg: ExportOptions): string
}

const genericExport = async (
  keyType: 'private' | 'public',
  keyFormat: 'spki' | 'pkcs8',
  key: unknown,
) => {
  if (isKeyObject(key)) {
    if (key.type !== keyType) {
      throw new TypeError(`key is not a ${keyType} key`)
    }

    return (key as ExtractableKeyObject).export({ format: 'pem', type: keyFormat })
  }

  if (!isCryptoKey(key)) {
    throw new TypeError(invalidKeyInput(key, 'CryptoKey', 'KeyObject'))
  }

  if (!key.extractable) {
    throw new TypeError('CryptoKey is not extractable')
  }

  if (key.type !== keyType) {
    throw new TypeError(`key is not a ${keyType} key`)
  }

  return formatPEM(
    encodeBase64(new Uint8Array(await crypto.subtle.exportKey(keyFormat, key))),
    `${keyType.toUpperCase()} KEY`,
  )
}

export const toSPKI = (key: unknown): Promise<string> => genericExport('public', 'spki', key)

export const toPKCS8 = (key: unknown): Promise<string> => genericExport('private', 'pkcs8', key)

/** Helper function to compare two byte arrays for equality */
const bytesEqual = (a: Uint8Array, b: readonly number[]): boolean => {
  if (a.byteLength !== b.length) return false
  for (let i = 0; i < a.byteLength; i++) {
    if (a[i] !== b[i]) return false
  }
  return true
}

/** ASN.1 DER parsing state */
interface ASN1State {
  readonly data: Uint8Array
  pos: number
}

/** Creates ASN.1 parsing state */
const createASN1State = (data: Uint8Array): ASN1State => ({ data, pos: 0 })

/** Reads the byte at the current position, failing fast when the input is truncated */
const readByte = (state: ASN1State): number => {
  const byte = state.data[state.pos++]
  if (byte === undefined) {
    throw new Error('Unexpected end of ASN.1 input')
  }
  return byte
}

/** Parses ASN.1 length encoding (both short and long form) */
const parseLength = (state: ASN1State): number => {
  const first = readByte(state)
  if (first & 0x80) {
    // Long form: first byte indicates number of subsequent length bytes
    const lengthOfLen = first & 0x7f
    let length = 0
    for (let i = 0; i < lengthOfLen; i++) {
      length = (length << 8) | readByte(state)
    }
    return length
  }
  // Short form: length is encoded directly in first byte
  return first
}

/** Skips ASN.1 elements (tag + length + content) */
const skipElement = (state: ASN1State, count: number = 1): void => {
  while (count-- > 0) {
    state.pos++ // Skip tag byte
    const length = parseLength(state)
    state.pos += length // Skip content bytes
  }
}

/** Expects a specific tag and throws if not found */
const expectTag = (state: ASN1State, expectedTag: number, errorMessage: string): void => {
  if (readByte(state) !== expectedTag) {
    throw new Error(errorMessage)
  }
}

/** Gets a subarray from current position */
const getSubarray = (state: ASN1State, length: number): Uint8Array => {
  if (length < 0 || state.pos + length > state.data.length) {
    throw new Error('Unexpected end of ASN.1 input')
  }
  const result = state.data.subarray(state.pos, state.pos + length)
  state.pos += length
  return result
}

/** Parses algorithm OID and returns the OID bytes */
const parseAlgorithmOID = (state: ASN1State): Uint8Array => {
  expectTag(state, 0x06, 'Expected algorithm OID')
  const oidLen = parseLength(state)
  return getSubarray(state, oidLen)
}

/** Parses a PKCS#8 or SPKI structure up to its algorithm identifier. */
function parseKeyHeader(state: ASN1State, keyFormat: 'spki' | 'pkcs8') {
  expectTag(state, 0x30, `Invalid ${keyFormat === 'spki' ? 'SPKI' : 'PKCS#8'} structure`)
  parseLength(state) // Skip outer length

  if (keyFormat === 'pkcs8') {
    // Skip version (INTEGER)
    expectTag(state, 0x02, 'Expected version field')
    const length = parseLength(state)
    state.pos += length
  }

  // Parse AlgorithmIdentifier SEQUENCE
  expectTag(state, 0x30, 'Expected algorithm identifier')
  parseLength(state)
}

/** Parses algorithm identifier and returns curve name for EC/ECDH keys */
const parseECAlgorithmIdentifier = (state: ASN1State): string => {
  const algOid = parseAlgorithmOID(state)

  // id-x25519
  if (bytesEqual(algOid, [0x2b, 0x65, 0x6e])) {
    return 'X25519'
  }

  // id-ecPublicKey 1.2.840.10045.2.1
  if (!bytesEqual(algOid, [0x2a, 0x86, 0x48, 0xce, 0x3d, 0x02, 0x01])) {
    throw new Error('Unsupported key algorithm')
  }

  // Parse curve parameters (should be an OID for named curves)
  expectTag(state, 0x06, 'Expected curve OID')
  const curveOidLen = parseLength(state)
  const curveOid = getSubarray(state, curveOidLen)

  if (bytesEqual(curveOid, [0x2a, 0x86, 0x48, 0xce, 0x3d, 0x03, 0x01, 0x07])) return 'P-256'
  if (bytesEqual(curveOid, [0x2b, 0x81, 0x04, 0x00, 0x22])) return 'P-384'
  if (bytesEqual(curveOid, [0x2b, 0x81, 0x04, 0x00, 0x23])) return 'P-521'

  throw new Error('Unsupported named curve')
}

const genericImport = async (
  keyFormat: 'spki' | 'pkcs8',
  keyData: Uint8Array,
  alg: string,
  options?: KeyImportOptions,
) => {
  const extractable = validateExtractableOption(options?.extractable)
  const entry = keyAlgorithm(alg, algArgument)
  if (entry.secret) {
    unsupportedAlg(algArgument)
  }
  const isPublic = keyFormat === 'spki'

  let algorithm: { name: string; hash?: string; namedCurve?: string }
  if (entry.resolve) {
    try {
      const state = createASN1State(keyData)
      parseKeyHeader(state, keyFormat)
      algorithm = entry.resolve({ crv: parseECAlgorithmIdentifier(state) })
    } catch {
      throw new JOSENotSupported('Invalid or unsupported key format')
    }
  } else {
    algorithm = entry.subtle
  }

  return crypto.subtle.importKey(
    keyFormat,
    keyData as Uint8Array<ArrayBuffer>,
    algorithm,
    extractable ?? isPublic,
    entry.usages[isPublic ? 0 : 1],
  )
}

type PEMImportFunction = (
  pem: string,
  alg: string,
  options?: KeyImportOptions,
) => Promise<types.CryptoKey>

/** Helper function to process PEM-encoded data */
const processPEMData = (pem: string, pattern: RegExp): Uint8Array => {
  return decodeBase64(pem.replace(pattern, ''))
}

export const fromPKCS8: PEMImportFunction = (pem, alg, options?) => {
  const keyData = processPEMData(pem, /(?:-----(?:BEGIN|END) PRIVATE KEY-----|\s)/g)
  return genericImport('pkcs8', keyData, alg, options)
}

export const fromSPKI: PEMImportFunction = (pem, alg, options?) => {
  const keyData = processPEMData(pem, /(?:-----(?:BEGIN|END) PUBLIC KEY-----|\s)/g)
  return genericImport('spki', keyData, alg, options)
}

/**
 * Extracts the Subject Public Key Info (SPKI) from an X.509 certificate. Parses the ASN.1 DER
 * structure to locate and extract the public key portion.
 *
 * @param buf - DER-encoded X.509 certificate bytes
 *
 * @returns SPKI structure as bytes
 */
function spkiFromX509(buf: Uint8Array): Uint8Array {
  const state = createASN1State(buf)

  // Parse outer certificate SEQUENCE
  expectTag(state, 0x30, 'Invalid certificate structure')
  const certificateLength = parseLength(state)
  if (certificateLength < 0 || state.pos + certificateLength > state.data.length) {
    throw new Error('Unexpected end of ASN.1 input')
  }

  // Parse tbsCertificate (To Be Signed Certificate) SEQUENCE
  expectTag(state, 0x30, 'Invalid tbsCertificate structure')
  parseLength(state) // Skip tbsCertificate length

  if (buf[state.pos] === 0xa0) {
    // Optional version field present (context-specific [0])
    // Skip: version, serialNumber, signature algorithm, issuer, validity, subject
    skipElement(state, 6)
  } else {
    // No version field (defaults to v1)
    // Skip: serialNumber, signature algorithm, issuer, validity, subject
    skipElement(state, 5)
  }

  // Extract subjectPublicKeyInfo SEQUENCE
  const spkiStart = state.pos
  expectTag(state, 0x30, 'Invalid SPKI structure')
  const spkiContentLen = parseLength(state)

  // Return the complete SPKI structure (tag + length + content)
  return buf.subarray(spkiStart, spkiStart + spkiContentLen + (state.pos - spkiStart))
}

/**
 * Extracts SPKI from a PEM-encoded X.509 certificate string.
 *
 * @param x509 - PEM-encoded X.509 certificate
 *
 * @returns SPKI structure as bytes
 */
export const fromX509: PEMImportFunction = (pem, alg, options?) => {
  let spki: Uint8Array
  try {
    const certificate = processPEMData(pem, /(?:-----(?:BEGIN|END) CERTIFICATE-----|\s)/g)
    spki = spkiFromX509(certificate)
  } catch (cause) {
    throw new TypeError('Failed to parse the X.509 certificate', { cause })
  }
  return genericImport('spki', spki, alg, options)
}
