import { PublicKey, SIGNATURE_FLAG_TO_SCHEME, SignatureFlag, SignatureScheme } from '@mysten/sui.js/cryptography';
import { Ed25519PublicKey } from '@mysten/sui.js/keypairs/ed25519';
import { Secp256k1PublicKey } from '@mysten/sui.js/keypairs/secp256k1';
import { Secp256r1PublicKey } from '@mysten/sui.js/keypairs/secp256r1';
import { fromB64 } from '@mysten/sui.js/utils';
import { ZkLoginPublicIdentifier } from '@mysten/sui.js/zklogin';

/**
 * Class representing a Public Key Serializer/Deserializer.
 */
export class PublicKeySerde {
  /**
   * Serializes a public key.
   *
   * @param {PublicKey} publicKey - The public key to serialize.
   * @return {Object} An object containing the serialized public key and the signature scheme.
   *   - publicKeyEncoded: The serialized public key encoded in base64.
   *   - scheme: The signature scheme corresponding to the public key flag.
   */
  static ser(publicKey: PublicKey): { publicKeyEncoded: string; schema: SignatureScheme } {
    return {
      publicKeyEncoded: publicKey.toBase64(),
      schema: SIGNATURE_FLAG_TO_SCHEME[publicKey.flag() as SignatureFlag],
    };
  }

  /**
   * Generates a public key object based on the provided input.
   *
   * @param {object} input - The input object containing the public key and the signature scheme.
   * @param {string | Uint8Array} input.publicKey - The public key.
   * @param {SignatureScheme} input.schema - The signature scheme.
   *
   * @return {PublicKey} - The generated public key object.
   *
   * @throws {Error} - If the signature scheme is unsupported.
   */
  static de(input: { publicKeyEncoded: string | Uint8Array; schema: string }): PublicKey {
    let pkEncoded =
      typeof input.publicKeyEncoded === 'string' ? fromB64(input.publicKeyEncoded) : input.publicKeyEncoded;

    let { schema } = input;
    if (pkEncoded.length === 34 || (input.schema === 'ED25519' && pkEncoded.length === 33)) {
      if (pkEncoded[0] === 0x00 || pkEncoded[0] === 0x01 || pkEncoded[0] === 0x02) {
        schema = SIGNATURE_FLAG_TO_SCHEME[pkEncoded[0]];
        pkEncoded = pkEncoded.slice(1);
      }
    }

    switch (schema) {
      case 'ED25519':
        return new Ed25519PublicKey(pkEncoded);
      case 'Secp256k1':
        return new Secp256k1PublicKey(pkEncoded);
      case 'Secp256r1':
        return new Secp256r1PublicKey(pkEncoded);
      case 'ZkLogin':
        return new ZkLoginPublicIdentifier(pkEncoded);
      default:
        throw new Error('Unsupported signature scheme: $input.scheme');
    }
  }
}