/**
 * @packageDocumentation
 *
 * Connection protection management for libp2p leveraging PSK encryption via XSalsa20.
 *
 * @example
 *
 * ```typescript
 * import { createLibp2p } from 'libp2p'
 * import { preSharedKey, generateKey } from 'libp2p/pnet'
 *
 * // Create a Uint8Array and write the swarm key to it
 * const swarmKey = new Uint8Array(95)
 * generateKey(swarmKey)
 *
 * const node = await createLibp2p({
 *   // ...other options
 *   connectionProtector: preSharedKey({
 *     psk: swarmKey
 *   })
 * })
 * ```
 */

import { randomBytes } from '@libp2p/crypto'
import { CodeError } from '@libp2p/interface/errors'
import { logger } from '@libp2p/logger'
import { handshake } from 'it-handshake'
import map from 'it-map'
import { duplexPair } from 'it-pair/duplex'
import { pipe } from 'it-pipe'
import { codes } from '../errors.js'
import {
  createBoxStream,
  createUnboxStream,
  decodeV1PSK
} from './crypto.js'
import * as Errors from './errors.js'
import { NONCE_LENGTH } from './key-generator.js'
import type { ConnectionProtector, MultiaddrConnection } from '@libp2p/interface/connection'

const log = logger('libp2p:pnet')

export { generateKey } from './key-generator.js'

export interface ProtectorInit {
  enabled?: boolean
  psk: Uint8Array
}

class PreSharedKeyConnectionProtector implements ConnectionProtector {
  public tag: string
  private readonly psk: Uint8Array
  private readonly enabled: boolean

  /**
   * Takes a Private Shared Key (psk) and provides a `protect` method
   * for wrapping existing connections in a private encryption stream.
   */
  constructor (init: ProtectorInit) {
    this.enabled = init.enabled !== false

    if (this.enabled) {
      const decodedPSK = decodeV1PSK(init.psk)
      this.psk = decodedPSK.psk
      this.tag = decodedPSK.tag ?? ''
    } else {
      this.psk = new Uint8Array()
      this.tag = ''
    }
  }

  /**
   * Takes a given Connection and creates a private encryption stream
   * between its two peers from the PSK the Protector instance was
   * created with.
   */
  async protect (connection: MultiaddrConnection): Promise<MultiaddrConnection> {
    if (!this.enabled) {
      return connection
    }

    if (connection == null) {
      throw new CodeError(Errors.NO_HANDSHAKE_CONNECTION, codes.ERR_INVALID_PARAMETERS)
    }

    // Exchange nonces
    log('protecting the connection')
    const localNonce = randomBytes(NONCE_LENGTH)

    const shake = handshake(connection)
    shake.write(localNonce)

    const result = await shake.reader.next(NONCE_LENGTH)

    if (result.value == null) {
      throw new CodeError(Errors.STREAM_ENDED, codes.ERR_INVALID_PARAMETERS)
    }

    const remoteNonce = result.value.slice()
    shake.rest()

    // Create the boxing/unboxing pipe
    log('exchanged nonces')
    const [internal, external] = duplexPair<Uint8Array>()
    pipe(
      external,
      // Encrypt all outbound traffic
      createBoxStream(localNonce, this.psk),
      shake.stream,
      (source) => map(source, (buf) => buf.subarray()),
      // Decrypt all inbound traffic
      createUnboxStream(remoteNonce, this.psk),
      external
    ).catch(log.error)

    return {
      ...connection,
      ...internal
    }
  }
}

export function preSharedKey (init: ProtectorInit): () => ConnectionProtector {
  return () => new PreSharedKeyConnectionProtector(init)
}
