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

import { type BaseStore, Activities, StoreInMemory } from '../storage/index.js';

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

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; passphrase?: 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,
    deps: WalletDependencies,
    genesisData: WalletHDGenesisData,
  ) {
    // Assign defaults if not otherwise provideed.
    const activities = deps.activities
      ? deps.activities
      : await Activities.from();
    const cache = deps.cache ? deps.cache : StoreInMemory.from();

    // Assign our final dependencies.
    const depsFinal = {
      activities,
      cache,
      blockchain: deps.blockchain,
    };

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

  // Dependencies.
  public readonly activities: Activities;
  public readonly blockchain: WalletBlockchain;
  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(genesisData.seed);
    } else if ('mnemonic' in genesisData) {
      // TODO: We really want to support multiple word-lists somehow.
      //       But this is not trivial because other word-lists are technically not standard.
      //       For now, we are allowing any phrase (BIP39 derives seed from literal phrase, not word indexes).
      const mnemonic = Mnemonic.fromRaw({
        words: genesisData.mnemonic.split(' '),
        passphrase: genesisData.passphrase,
      });

      // Create the HD Node from our Mnemonic instance.
      this.hdPrivateNode = HDPrivateNode.fromMnemonic(mnemonic);
    } else {
      this.hdPrivateNode = HDPrivateNode.fromXPriv(genesisData.xpriv);
    }

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

  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 = ExtMap.flatten(unspentDirectives);

    // 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<TransactionBuilder> {
    // 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.
    return TransactionBuilder.build(
      txTemplate,
      unspents.toArray(),
      changeLockingBytecode,
    );
  }

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

    // Convert the tx to bytes.
    const generatedTxBytes = generatedTx.toBytes();

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

    // Return the transaction.
    return 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(
        {
          activities: this.activities,
          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;
}