import { PaginatedTransactionResponse, SuiClient } from '@mysten/sui.js/client';
import {
  PublicKey,
  SIGNATURE_FLAG_TO_SCHEME,
  SerializedSignature,
  SignatureFlag,
  SignatureScheme,
  parseSerializedSignature,
} 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 { MultiSigPublicKey } from '@mysten/sui.js/multisig';
import { verifyPersonalMessage, verifyTransactionBlock } from '@mysten/sui.js/verify';

import { PublicKeySerde } from '@/PublicKeySerde';
import { MessageType, PublicKeyWithSchema, SuiAddress } from '@/types';
import { isSameAddress, stringToBuffer } from '@/utils';

/**
 * Get public key from signature
 * @param message message
 * @param type message type
 * @param signature signature
 * @returns public key
 */
export async function getPublicKeyFromSignature(
  message: Uint8Array,
  type: MessageType,
  signature: SerializedSignature,
): Promise<PublicKey> {
  switch (type) {
    case 'TransactionBlock':
      return verifyTransactionBlock(message, signature);
    case 'Personal':
      return verifyPersonalMessage(message, signature); // options);
    default:
      throw new Error('Invalid message type');
  }
}

/**
 * Get public key from personal signature
 * @param message message
 * @param signature  signature
 * @returns public key
 */
export async function getPublicKeyFromPersonalSignature(
  message: string,
  signature: SerializedSignature,
): Promise<PublicKey> {
  return getPublicKeyFromSignature(stringToBuffer(message), 'Personal', signature);
}

/**
 * Verify signature
 * @param message message
 * @param type message type
 * @param signature signature
 * @param address target address
 */
export async function verifySignature(
  message: Uint8Array,
  type: MessageType,
  signature: SerializedSignature,
  address: SuiAddress,
): Promise<boolean> {
  try {
    const publicKey = await getPublicKeyFromSignature(message, type, signature);
    return isSameAddress(publicKey.toSuiAddress(), address);
  } catch {
    return false;
  }
}

/**
 * Verify personal signature
 * @param message message
 * @param signature signatured
 * @param address target address
 * @returns
 */
export async function verifyPersonalSignature(
  message: string,
  signature: SerializedSignature,
  address: SuiAddress,
): Promise<boolean> {
  return verifySignature(stringToBuffer(message), 'Personal', signature, address);
}

/**
 * Verify transaction signature
 * @param payload transaction payload
 * @param signature signature
 * @param address target address
 */
export async function verifyTransactionSignature(
  payload: Uint8Array,
  signature: SerializedSignature,
  address: SuiAddress,
): Promise<boolean> {
  return verifySignature(payload, 'TransactionBlock', signature, address);
}

/**
 * Serialize public key
 * @Deprecated: Use PublicKeySerde.ser
 * @param publicKey public key object
 * @returns public key with scheme
 */
export function serializePublicKey(publicKey: PublicKey): PublicKeyWithSchema {
  return {
    publicKeyEncoded: publicKey.toBase64(),
    schema: SIGNATURE_FLAG_TO_SCHEME[publicKey.flag() as SignatureFlag],
  };
}

/**
 * Deserialize public key
 * @Deprecated: Use PublicKeySerde.de
 * @param publicKey public key string
 * @param scheme public key scheme
 * @returns public key object
 */
export function deserializePublicKey(publicKey: string | Uint8Array, scheme: SignatureScheme): PublicKey {
  switch (scheme) {
    case 'ED25519':
      return new Ed25519PublicKey(publicKey);
    case 'Secp256k1':
      return new Secp256k1PublicKey(publicKey);
    case 'Secp256r1':
      return new Secp256r1PublicKey(publicKey);
    default:
      throw new Error('Unsupported signature scheme');
  }
}

export class PublicKeyManager {
  private cachedKeys: Map<SuiAddress, PublicKey>;

  constructor(private readonly client: SuiClient) {
    this.cachedKeys = new Map<SuiAddress, PublicKey>();
  }

  /**
   * Get address public key from cache or from chain
   * @param address target address
   */
  async getPublicKey(address: SuiAddress) {
    const cachedKey = this.cachedKeys.get(address);

    if (cachedKey) {
      return cachedKey;
    }

    const publicKey = await this.fetchPublicKey(address);
    if (publicKey) {
      this.cachedKeys.set(address, publicKey);
    }
    return publicKey;
  }

  /**
   * Batch get public keys from addresses
   * @param addresses addresses
   */
  async getPublicKeys(addresses: SuiAddress[]) {
    const results: (PublicKey | undefined)[] = new Array(addresses.length).fill(undefined);

    for (let index = 0; index < addresses.length; index++) {
      const address = addresses[index];
      results[index] = this.cachedKeys.get(address);
    }

    for (let index = 0; index < results.length; index++) {
      const result = results[index];
      if (result === undefined) {
        results[index] = await this.fetchPublicKey(addresses[index]);
      }
    }

    for (let index = 0; index < addresses.length; index++) {
      const address = addresses[index];
      if (results[index]) {
        this.cachedKeys.set(address, results[index] as PublicKey);
      }
    }

    return results;
  }

  private async fetchPublicKey(address: SuiAddress) {
    const publicKeyFromChain = await this.fetchPublicKeyFromChain(address);
    if (publicKeyFromChain) {
      return publicKeyFromChain;
    }
    return undefined;
  }

  private async fetchPublicKeyFromChain(address: SerializedSignature) {
    let txs: PaginatedTransactionResponse;
    try {
      txs = await this.client.queryTransactionBlocks({
        filter: { FromAddress: address },
        options: { showInput: true },
        limit: 2,
      });
    } catch {
      return undefined;
    }

    if (txs.data.length === 0 || !txs.data[0].transaction?.txSignatures) {
      return undefined;
    }

    const tx = txs.data[0];

    if (!tx.transaction) {
      return undefined;
    }

    const signatures = tx.transaction.txSignatures as string[];

    for (let index = 0; index < signatures.length; index++) {
      const signature = signatures[index];
      const publicKey = this.getAddressFromSignature(signature, address);
      if (!publicKey) {
        continue;
      }
      return publicKey;
    }
    return undefined;
  }

  private getAddressFromSignature(signature: SerializedSignature, address: SuiAddress) {
    const parsedSignature = parseSerializedSignature(signature);

    switch (parsedSignature.signatureScheme) {
      case 'MultiSig': {
        const multisigAddress = new MultiSigPublicKey(parsedSignature.multisig.multisig_pk).toSuiAddress();
        if (isSameAddress(multisigAddress, address)) {
          throw new Error('multisig wallet cannot be owner');
        }
        return undefined;
      }
      case 'ZkLogin':
        if (isSameAddress(parsedSignature.zkLogin.address, address)) {
          throw new Error('ZkLogin wallet cannot be owner');
        }
        return undefined;
      case 'ED25519':
      case 'Secp256k1':
      case 'Secp256r1': {
        const publicKey = PublicKeySerde.de({
          publicKeyEncoded: parsedSignature.publicKey,
          schema: parsedSignature.signatureScheme,
        });
        if (isSameAddress(publicKey.toSuiAddress(), address)) {
          return publicKey;
        }
        return undefined;
      }
      default:
        throw new Error('Invalid signature scheme');
    }
  }
}