import { mnemonicToSeedSync } from "@scure/bip39";
import { HDPrivateKey } from "libnexa-ts";
import type { AccountType } from "../utils/enums";
import type { KeyPath } from "../types/wallet.types";
import { keyPathToString, stringToKeyPath } from "../utils/keypath";

export class KeyManager {

  private seed!: Uint8Array;
  private masterKey!: HDPrivateKey;
  private accountKeys = new Map<number, HDPrivateKey>();
  private walletKeys = new Map<string, HDPrivateKey>();

  public init(mnemonic: string | Uint8Array): void {
    this.seed = typeof mnemonic === "string" ? mnemonicToSeedSync(mnemonic) : mnemonic;
    this.masterKey = HDPrivateKey.fromSeed(this.seed).deriveChild(44, true).deriveChild(29223, true);
  }

  public reset(): void {
    if (!this.seed) {
      throw new Error("KeysManager not initialized");
    }
    this.masterKey = HDPrivateKey.fromSeed(this.seed).deriveChild(44, true).deriveChild(29223, true);
    this.accountKeys.clear();
    this.walletKeys.clear();
  }

  private getAccountKey(account: AccountType): HDPrivateKey {
    let key = this.accountKeys.get(account);
    if (!key) {
      key = this.masterKey.deriveChild(account, true);
      this.accountKeys.set(account, key);
    }

    return key;
  }

  public getKey(keyPath: string | KeyPath): HDPrivateKey {
    const path = typeof keyPath === 'string'
      ? keyPath
      : keyPathToString(keyPath.account, keyPath.type, keyPath.index);

    let key = this.walletKeys.get(path);
    if (!key) {
      const { account, type, index } = typeof keyPath === 'string'
        ? stringToKeyPath(keyPath)
        : keyPath;
      const accountKey = this.getAccountKey(account);
      key = accountKey.deriveChild(type, false).deriveChild(index, false);
      this.walletKeys.set(path, key);
    }

    return key;
  }
}