import { BaseWallet } from './base-wallet.js';
import { type TransactionTemplate } from './types.js';
import { TransactionBuilder } from './transaction-builder.js';
import { type WalletDependencies } from './types.js';
import { WalletP2PKH } from './wallet-p2pkh.js';

import { type BaseBlockchain } from '../blockchain';
import { type BaseStore, Activities, StoreInMemory } from '../storage';

import {
  type TokenBalances,
  HDPrivateNode,
  ExtMap,
  Mnemonic,
  calculateBalanceSats,
  calculateBalanceTokens,
  encodeExtendedJson,
} from '@xocash/primitives';

import { hash256, binToHex, hexToBin, utf8ToBin } from '@bitauth/libauth';

export const DERIVATION_PATH = "m/44'/145'/0'";
export const ADDRESS_GAP = 20;

// These constants correspond to the "change" path in BIP44.
// See: https://github.com/bitcoin/bips/blob/master/bip-0044.mediawiki
export const CHAIN_EXTERNAL = 0;
export const CHAIN_INTERNAL = 1;

export type WalletHDGenesisData = {
  type: 'WalletHD';
  derivationPath: string;
  addressGap?: number;
} & ({ seed: string } | { mnemonic: string } | { xpriv: string });

/**
 * A Hierarhical Determinstic (BIP44) Wallet.
 *
 * @example
 *
 * ```ts
 * // Define services to use.
 * const services = {
 *   blockchain,
 *   cache,
 *   storage,
 * }
 *
 * // Instantiate an instance of WalletHD.
 * const walletHD = await WalletHD.from({
 *   derivationPath: "m/44'/145'/0'",
 *   mnemonic: 'page pencil stock planet limb cluster assault speak off joke private pioneer'
 *   // or "seed" or "xpub"
 * }, services);
 *
 * // Scan for transactions and start monitoring.
 * await wallet = await walletHD.start();
 * ```
 */
export class WalletHD extends BaseWallet {
  static async from<T extends typeof WalletHD>(
    this: T,
    services: WalletDependencies,
    genesisData: WalletHDGenesisData,
  ) {
    // Assign in memory stores as defaults in case otherwise not provided.
    const servicesFinal = {
      ...{
        // Default to an in-memory-only Activity store.
        activities: await Activities.fromStores(StoreInMemory.from()),
        store: new StoreInMemory(),
        cache: new StoreInMemory(),
      },
      ...services,
    };

    // Return a new instance of WalletHD.
    return new this(servicesFinal, genesisData) as InstanceType<T>;
  }

  // Dependencies.
  public readonly activities: Activities;
  public readonly blockchain: BaseBlockchain;
  public readonly cache: BaseStore;

  // Genesis Data.
  public readonly genesisData: WalletHDGenesisData;

  // Optimizations/Simplifications.
  public readonly hdPrivateNode: HDPrivateNode;

  // Mutable State.
  public shouldMonitor = false;

  // Store child wallets, grouped by chain ("change") path in BIP44 spec.
  public wallets: { [chainPath: number]: Array<WalletP2PKH> } = {};

  constructor(
    dependencies: Required<WalletDependencies>,
    genesisData: WalletHDGenesisData,
  ) {
    super();

    // Assign our dependencies.
    this.activities = dependencies.activities;
    this.blockchain = dependencies.blockchain;
    this.cache = dependencies.cache;

    // Assign our state.
    this.genesisData = genesisData;

    // Instantiate the HDPrivateNode from the type specified in our genesis data.
    if ('seed' in genesisData) {
      this.hdPrivateNode = HDPrivateNode.fromSeed(hexToBin(genesisData.seed));
    } else if ('mnemonic' in genesisData) {
      this.hdPrivateNode = HDPrivateNode.fromMnemonic(
        Mnemonic.fromPhrase(genesisData.mnemonic),
      );
    } else {
      this.hdPrivateNode = HDPrivateNode.fromXPriv(genesisData.xpriv);
    }

    // Initialize empty arrays for our wallets.
    this.wallets[CHAIN_EXTERNAL] = [];
    this.wallets[CHAIN_INTERNAL] = [];
  }

  async getId() {
    // NOTE: Use the cache to optimize future calls.
    return binToHex(
      await this.cache.getOrSet('walletId', async () =>
        this.deriveKey('walletId'),
      ),
    );
  }

  async deriveKey(keyName: string): Promise<Uint8Array> {
    // Hash the genesis data using Hash256 and the key ("WalletID") for this derivation using Hash256.
    const genesisDataHashed = hash256(
      utf8ToBin(encodeExtendedJson(this.genesisData)),
    );
    const keyHashed = hash256(utf8ToBin(keyName));

    // Combine the hashes and do one final hash for good measure.
    return hash256(new Uint8Array([...genesisDataHashed, ...keyHashed]));
  }

  async getEncryptionKey() {
    // NOTE: Use the cache to optimize future calls.
    return await this.cache.getOrSet('encryptionKey', async () =>
      this.deriveKey('encryptionKey'),
    );
  }

  async getSigningKey() {
    // NOTE: Use the cache to optimize future calls.
    return this.cache.getOrSet('signingKey', async () =>
      this.deriveKey('signingKey'),
    );
  }

  async getActivities(): Promise<Activities> {
    return this.activities;
  }

  async start() {
    // Get the address gap.
    // TODO: This should come from activities.
    const addressGap = this.genesisData.addressGap || ADDRESS_GAP;

    // Scan the HDWallet for child wallets.
    await this.scan(0, addressGap);

    // Declare storage for our start promises.
    const startPromises: Promise<void>[] = [];

    // Iterate over each chain path and start each wallet.
    for (const chainWallets of Object.values(this.wallets)) {
      for (const wallet of chainWallets) {
        startPromises.push(wallet.start());
      }
    }

    // Wait for all of our start promises to complete.
    await Promise.all(startPromises);
  }

  async stop() {
    // Declare storage for our stop promises.
    const stopPromises: Promise<void>[] = [];

    // Iterate over each chain path and stop each wallet.
    for (const chainWallets of Object.values(this.wallets)) {
      for (const wallet of chainWallets) {
        stopPromises.push(wallet.stop());
      }
    }

    // Wait for all of our stop promises to complete.
    await Promise.all(stopPromises);
  }

  async destroy() {
    this.removeAllListeners();
  }

  async getTransactions() {
    // Declare storage for our fetch promises.
    const fetchPromises = [];

    // Iterate over each chain path and fetch transactions for each wallet.
    for (const chainWallets of Object.values(this.wallets)) {
      for (const wallet of chainWallets) {
        fetchPromises.push(wallet.getTransactions());
      }
    }

    // Wait for all of our fetch promises to complete.
    const txs = await Promise.all(fetchPromises);

    // Flatten (and deduplicate) the transactions.
    const txsFlattened = ExtMap.flatten(txs);

    // Emit an event to notify that transactions have been updated.
    this.emit('transactionsUpdated', txsFlattened);

    // Return the deduplicated transactions.
    return txsFlattened;
  }

  async getUnspents() {
    // Declare storage for our fetch promises.
    const fetchPromises = [];

    // Iterate over each chain path and fetch unspents for each wallet.
    for (const chainWallets of Object.values(this.wallets)) {
      for (const wallet of chainWallets) {
        fetchPromises.push(wallet.getUnspents());
      }
    }

    // Wait for all of our fetch promises to complete.
    const unspents = await Promise.all(fetchPromises);

    // Flatten them.
    const unspentsFlattened = ExtMap.flatten(unspents);

    // Return the unspents.
    return unspentsFlattened;
  }

  async getUnspentDirectives() {
    // Declare storage for our fetch promises.
    const fetchPromises = [];

    // Iterate over each chain path and fetch unspent directives for each wallet.
    for (const chainWallets of Object.values(this.wallets)) {
      for (const wallet of chainWallets) {
        fetchPromises.push(wallet.getUnspentDirectives());
      }
    }

    // Wait for all of our fetch promises to complete.
    const unspentDirectives = await Promise.all(fetchPromises);

    // Flatten them.
    const unspentsFlattened = unspentDirectives.flat();

    // Return the unspents.
    return unspentsFlattened;
  }

  async getBalanceSats() {
    // Fetch our unspents.
    const unspents = await this.getUnspents();

    // Get the individual UTXOs.
    const utxos = unspents
      .map((blockchainUTXO) => blockchainUTXO.utxo)
      .toArray();

    // Calculate the balance.
    return calculateBalanceSats(utxos);
  }

  async getBalanceTokens(): Promise<TokenBalances> {
    // Fetch our unspents.
    const unspents = await this.getUnspents();

    // Get the individual UTXOs.
    const utxos = unspents
      .map((blockchainUTXO) => blockchainUTXO.utxo)
      .toArray();

    // Calculate the token balances.
    return calculateBalanceTokens(utxos);
  }

  async getAddresses() {
    const chainPaths = [CHAIN_EXTERNAL, CHAIN_INTERNAL];

    const addressPromises = [];

    for (const chainPath of chainPaths) {
      addressPromises.push(
        ...this.wallets[chainPath].map((wallet) =>
          wallet.getReceivingAddress(),
        ),
      );
    }

    const addresses = await Promise.all(addressPromises);

    return ExtMap.fromArray(addresses, (address) => address.toLockscriptHex());
  }

  async getReceivingAddress(chainPath = CHAIN_EXTERNAL) {
    // TODO: Derive a new wallet, not the first one you lazy fuckhead.
    //       You'll probably need to refactor your scan() approach.
    let receivingWallet = this.wallets[chainPath][0];

    // If there is no receiving wallet, create one.
    // TODO: This is fucked. You need to rethink deriveWallets.
    if (!receivingWallet) {
      this.deriveWallets(1, chainPath);

      receivingWallet = this.wallets[chainPath][0];
    }

    // Return the receiving address of the derived wallet.
    return receivingWallet.privateKey.derivePublicKey().deriveAddress();
  }

  async signTransaction(
    txTemplate: Partial<TransactionTemplate>,
  ): Promise<Uint8Array> {
    // Populate the final transaction template by merging with defaults.
    // NOTE: This is just a convenience to make "send" easier to use.
    const txTemplateFinal: TransactionTemplate = {
      ...{
        version: 2,
        locktime: 0,
        inputs: [],
        outputs: [],
      },
      ...txTemplate,
    };

    // Get the receiving address of this wallet.
    // NOTE: As per BIP44 spec, we set the change address to use the internal chain.
    const changeAddress = await this.getReceivingAddress(CHAIN_EXTERNAL);

    // Get the change locking bytecode (change address) for this wallet.
    const changeLockingBytecode = changeAddress.toLockscriptBytes();

    // Fetch unspents that are usable by this wallet.
    const unspents = await this.getUnspentDirectives();

    // Build the transaction.
    const tx = TransactionBuilder.build(
      txTemplateFinal,
      unspents,
      changeLockingBytecode,
    );

    // Generate the transaction.
    const generatedTxBytes = tx.toBytes();

    return generatedTxBytes;
  }

  async sendTransaction(txTemplate: Partial<TransactionTemplate>) {
    // Sign the transaction template.
    const generatedTxBytes = await this.signTransaction(txTemplate);

    // Broadcast the transaction to the blockchain.
    await this.blockchain.broadcastTransaction(generatedTxBytes);
  }

  async deriveWallets(
    count: number,
    chainPath = CHAIN_EXTERNAL,
    startIndex = 0,
  ) {
    // Create an array to store our wallets.
    const newWallets: Array<WalletP2PKH> = [];

    // Create the given count of wallets.
    for (let i = 0; i < count; i++) {
      // Define the derivation path.
      const derivationPath = `${this.genesisData.derivationPath}/${chainPath}/${startIndex + i}`;

      // Derive the private key at the given path.
      const privateKeyHex = await this.cache.getOrSet(derivationPath, () => {
        const childNode = this.hdPrivateNode.derivePath(derivationPath);
        return childNode.toPrivateKey().toHex();
      });

      // Create a child P2PKH wallet for this Private Key.
      const wallet = await WalletP2PKH.from(
        {
          blockchain: this.blockchain,
          cache: this.cache,
        },
        {
          type: 'WalletP2PKH',
          privateKey: privateKeyHex,
        },
      );

      // Subscribe to wallet-state updates.
      wallet.on('stateUpdated', async () => {
        this.emit('stateUpdated', null);
      });

      // Start monitoring the wallet.
      wallet.start();

      // Add the wallet to our list of wallets.
      newWallets.push(wallet);
    }

    // Set our wallets.
    this.wallets[chainPath] = [...this.wallets[chainPath], ...newWallets];

    // Return the wallets.
    return newWallets;
  }

  async scan(
    startIndex = 0,
    addressGap = 20,
    chainPaths: Array<number> = [CHAIN_EXTERNAL, CHAIN_INTERNAL],
  ) {
    // Iterate over each of the provided chain (change) paths.
    for (const chainPath of chainPaths) {
      // Array to store active wallet nodes
      const wallets: Array<WalletP2PKH> = [];

      let currentIndex = startIndex;
      let emptyAddressCount = 0;

      // Continue scanning until we find 'addressGap' consecutive empty addresses
      while (emptyAddressCount < addressGap) {
        // Derive a number of wallets equivalent to our addressGap.
        const batch: Array<WalletP2PKH> = await this.deriveWallets(
          addressGap,
          chainPath,
          currentIndex,
        );

        // Get the transaction history of each wallet in the batch
        const results = await Promise.all(
          batch.map((wallet) => wallet.getTransactions()),
        );

        // Process the results
        for (let i = 0; i < results.length; i++) {
          if (results[i].size > 0) {
            // This address has transactions, add it to nodes and reset empty address count
            wallets.push(batch[i]);
            emptyAddressCount = 0;
          } else {
            // This address is empty, increment the empty address count
            emptyAddressCount++;
          }

          currentIndex++;

          // If we've found 'addressGap' consecutive empty addresses, stop scanning
          if (emptyAddressCount >= addressGap) {
            break;
          }
        }
      }
    }

    this.emit('stateUpdated', null);
  }
}

//-----------------------------------------------------------------------------
// Factories Types/Functions
//-----------------------------------------------------------------------------

export type WalletHDFactory<T extends WalletHD = WalletHD> = (
  services: WalletDependencies,
  genesisData: WalletHDGenesisData,
) => Promise<T>;

export async function useWalletHD<T extends WalletHD = WalletHD>(
  services: WalletDependencies,
  genesisData: WalletHDGenesisData,
): Promise<T> {
  return (await WalletHD.from(services, genesisData)) as T;
}