import * as utils from '../../utils.js';
import { SupportedKeyDerivationTypes } from '../module.js';
import {
  DerivationNonce,
  IDerivedKey,
  IEncryptableKey,
  IKey,
  IKeyDerivationStrategy,
  IPublicKey,
  KeyInformation,
} from './interfaces.js';
import { jose, crypto } from '../../utils.js';

/**
 * Implementation of a key derivation strategy for AES-256 GCM keys.
 */
export class AES256GCMKeyDerivationStrategy implements IKeyDerivationStrategy {
  /**
   * Recover a derived key from a master key and a derivation nonce.
   * @param masterKey The master key used for derivation.
   * @param nonce The derivation nonce used to derive the key.
   * @returns A Promise that resolves to an AES256GCMKey.
   */
  hydrate(masterKey: CryptoKey, nonce: DerivationNonce): Promise<AES256GCMKey> {
    return this.create(masterKey, nonce.salt);
  }

  /**
   * Create a new key from a raw key represented as a Uint8Array.
   * @param rawKey The raw key as a Uint8Array.
   * @returns A Promise that resolves to an IKey implementation.
   */
  async fromRaw(rawKey: Uint8Array): Promise<IKey> {
    const key = await crypto.subtle.importKey('raw', rawKey, { name: 'AES-GCM' }, true, ['encrypt', 'decrypt']);
    return new AES256GCMKey(key);
  }

  /**
   * Create a new AES-256 GCM key derived from a master key and a salt value.
   * @param masterKey The master key for key derivation.
   * @param salt The salt value used for key derivation.
   * @returns A Promise that resolves to an AES256GCMKey.
   */
  async create(masterKey: CryptoKey, salt: string): Promise<AES256GCMKey> {
    const rawMasterKey = await crypto.subtle.exportKey('raw', masterKey);
    const pathBytes = new TextEncoder().encode(salt.toString());
    const derivedKey = await utils.pbkdf2(new Uint8Array(rawMasterKey), pathBytes, 256, 10000);
    const derivedCryptoKey = await crypto.subtle.importKey('raw', derivedKey, { name: 'AES-GCM' }, true, [
      'encrypt',
      'decrypt',
    ]);
    return new AES256GCMKey(derivedCryptoKey, undefined, new DerivationNonce(salt, 10000));
  }
}

// TODO: This isnt actually a public key. not sure what to do about this. do we need to?
export class AES256GCMFingerprint implements IPublicKey {
  constructor(public bytes: Uint8Array) {}

  getPublicKeyHex(): Promise<string> {
    return Promise.resolve(utils.hex.bufToHex(this.bytes));
  }
  getPublicKeyBytes(): Promise<Uint8Array> {
    return Promise.resolve(this.bytes);
  }

  verify(_signature: Uint8Array, _message: Uint8Array): Promise<boolean> {
    throw new Error('Method not implemented.');
  }
}

const AES_GCM_IV_LENGTH = 12;

export class AES256GCMKey implements IKey, IDerivedKey, IEncryptableKey {
  constructor(
    private key: CryptoKey,
    private rawSecretBytes?: Uint8Array,
    private derivation_path?: DerivationNonce,
  ) {}

  async toJWK({ includePrivate }: { includePrivate: boolean }) {
    const jwk = await jose.exportJWK(this.key);
    const thumbprint = await jose.calculateJwkThumbprint(jwk, 'sha256');
    if (!includePrivate) {
      delete jwk.d;
    }
    return {
      ...jwk,
      kid: thumbprint,
    };
  }

  async toCryptoKey(): Promise<CryptoKey> {
    return this.key;
  }

  async exportKeyInformation(): Promise<KeyInformation> {
    return {
      version: 1,
      curve: SupportedKeyDerivationTypes.AES256GCM,
      key: this.derivation_path
        ? {
            type: 'derived',
            from: 'master',
            nonce: this.derivation_path.toString(),
          }
        : {
            type: 'raw',
            rawKeyHex: utils.hex.bufToHex(this.rawSecretBytes!),
          },
      publicKey: utils.hex.bufToHex((await this.getPublicKey()).bytes),
      fingerprint: utils.hex.bufToHex(await this.getFingerprint()),
    };
  }

  getChainCode(): Uint8Array {
    throw new Error('Method not implemented.');
  }

  async encrypt(message: Uint8Array): Promise<Uint8Array> {
    const iv = crypto.getRandomValues(new Uint8Array(AES_GCM_IV_LENGTH));

    const encrypted = await crypto.subtle.encrypt({ name: 'AES-GCM', iv }, this.key, message);

    const concatCiphertext = new Uint8Array(iv.length + encrypted.byteLength);
    concatCiphertext.set(iv, 0);
    concatCiphertext.set(new Uint8Array(encrypted), AES_GCM_IV_LENGTH);

    return new Uint8Array(concatCiphertext);
  }

  async decrypt(encryptedMessage: Uint8Array): Promise<Uint8Array> {
    // First 12 bytes of ciphertext is IV
    const iv = encryptedMessage.slice(0, AES_GCM_IV_LENGTH);

    const decrypted = await crypto.subtle.decrypt(
      { name: 'AES-GCM', iv },
      this.key,
      encryptedMessage.slice(AES_GCM_IV_LENGTH),
    );

    return new Uint8Array(decrypted);
  }

  getPath(): DerivationNonce {
    return this.derivation_path!;
  }

  async getFingerprint(): Promise<Uint8Array> {
    const exportedKey = await crypto.subtle.exportKey('raw', this.key);
    const keyHash = await utils.sha256(new Uint8Array(exportedKey));

    return Promise.resolve(utils.toUint8Array(keyHash));
  }

  async getPublicKey(): Promise<AES256GCMFingerprint> {
    return new AES256GCMFingerprint(await this.getFingerprint());
  }
}
