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

import {
  type BaseBlockchain,
  type BlockchainTransactions,
  type BlockchainUTXOs,
} from '../blockchain';
import {
  type BaseStore,
  type StoreValueProxy,
  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 });

export type WalletHDEvents<WalletType> = {
  walletDataUpdated: WalletHDGenesisData;
  shouldMonitorUpdated: boolean;
  walletsUpdated: Array<WalletType>;
  transactionsUpdated: BlockchainTransactions;
  unspentsUpdated: BlockchainUTXOs;
  stateUpdated: string | null;
};

/**
 * 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<
  WalletType extends WalletP2PKH = WalletP2PKH,
> extends BaseWallet<WalletHDEvents<WalletType>> {
  // Mutable State.
  public shouldMonitor = false;

  // Store values that are proxied for convenience/clarity.
  // NOTE: We define these for clarity so that we can see the state stored with this wallet.
  public genesisData: StoreValueProxy<WalletHDGenesisData>;

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

  static async from<T extends typeof WalletHD>(
    this: T,
    services: WalletServices,
    genesisData?: WalletHDGenesisData,
    walletFactory?: WalletP2PKHFactory<any>,
  ) {
    // Assign in memory stores as defaults in case otherwise not provided.
    const servicesFinal = {
      ...{
        store: new StoreInMemory(),
        cache: new StoreInMemory(),
      },
      ...services,
    };

    // If we are attempting to create a new wallet...
    if (genesisData) {
      // Ensure that the provided store does not already contain data.
      if (Object.keys(await servicesFinal.store.all()).length) {
        throw new Error(`Provided store already contains data`);
      }

      // Store the wallet genesis data.
      await servicesFinal.store.set<WalletHDGenesisData>(
        'genesis',
        genesisData,
      );
    }

    // Retrieve our genesis data and extract the private key.
    const retrievedGenesisData =
      await servicesFinal.store.getOrFail<WalletHDGenesisData>('genesis');

    let hdPrivateNode: HDPrivateNode;

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

    // Return a new instance of WalletHD.
    return new this(
      servicesFinal.blockchain,
      servicesFinal.store,
      servicesFinal.cache,
      hdPrivateNode,
      walletFactory,
    ) as InstanceType<T>;
  }

  constructor(
    public blockchain: BaseBlockchain,
    public store: BaseStore,
    public cache: BaseStore,
    public hdPrivateNode: HDPrivateNode,
    private walletFactory: WalletP2PKHFactory<WalletType> = useWalletP2PKH,
  ) {
    super();

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

    // Initialize proxies to our stored data.
    this.genesisData = this.store.proxy('genesis');
  }

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

  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 start() {
    // Get the genesis data.
    const genesisData = await this.genesisData.getOrFail();

    // Get the address gap.
    const addressGap = 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];

    let 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();

    // Create an instance of our transaction builder and populate with our Wallet's UTXOs.
    // NOTE: We must deal with tokens first and satoshis last.
    //       The reason is that the change output will be determined based on transaction fee.
    const txBuilder = new TransactionBuilder(txTemplateFinal, unspents)
      .addTokensFromUnspents()
      .addTokenChangeOutputs(changeLockingBytecode)
      .addSatoshisFromUnspents()
      .addSatoshiChangeOutput(changeLockingBytecode);

    // Generate the transaction.
    const generatedTxBytes = txBuilder.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 addActivity<T>(activity: T) {
    // Generate an Activity ID.
    // TODO: Look more into how LibAuth intends to do this.
    const activityId = Date.now().toString();

    // Store the data.
    await this.store.set<T>(activityId, activity);
  }

  async export() {
    const allData = await this.store.all();

    const asJson = encodeExtendedJson(allData);

    return asJson;
  }

  async deriveWallets(
    count: number,
    chainPath = CHAIN_EXTERNAL,
    startIndex = 0,
  ) {
    // Get the derivation path from the genesis data.
    const genesisData = await this.genesisData.getOrFail();

    // Create an array to store our wallets.
    const newWallets: Array<WalletType> = [];

    // Create the given count of wallets.
    for (let i = 0; i < count; i++) {
      // Define the derivation path.
      const derivationPath = `${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 this.walletFactory(
        {
          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<WalletType> = [];

      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<WalletType> = 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);
  }

  protected async deriveKey(keyName: string): Promise<Uint8Array> {
    // Get the genesis data for this wallet.
    // TODO: This needs to change!
    const genesisData = await this.store.getOrFail('genesis');

    // Hash the genesis data using Hash256 and the key ("WalletID") for this derivation using Hash256.
    const genesisDataHashed = hash256(utf8ToBin(JSON.stringify(genesisData)));
    const keyHashed = hash256(utf8ToBin(keyName));

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

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

export type WalletHDFactory<T extends WalletHD = WalletHD> = (
  services: WalletServices,
  genesisData?: WalletHDGenesisData,
  walletFactory?: WalletP2PKHFactory<any>,
) => Promise<T>;

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