import type { KVStore, WalletDB } from "../persistence";
import type { Account, AccountDTO, AddressDTO, AddressKey, Balance, ITXHistory, RostrumParams, VaultDTO, VaultInfo } from "../types";
import { AccountType, keyPathToString, KeySpace, MAIN_WALLET_ID, sumBalance, sumTokensBalance } from "../utils";
import type { AssetService } from "./asset";
import { WalletDiscoveryService } from "./discovery";
import type { KeyManager } from "./key";
import type { RostrumService } from "./rostrum";
import type { SessionManager } from "./session";
import type { TransactionService } from "./transaction";

type WalletAddress = AccountDTO | AddressDTO | VaultDTO;

interface AddressHistory {
  address: WalletAddress;
  txs: ITXHistory[];
}

interface UpdateInfo {
  address: WalletAddress;
  result: string
}

export type WalletEvent =
  | { type: 'new_tip'; tip: number; }
  | { type: 'new_account'; account: Account; }
  | { type: 'new_vault'; vault: VaultInfo; }
  | { type: 'discover_wallet'; loading: boolean; }
  | { type: 'load_wallet'; loading: boolean; }
  | { type: 'sync_wallet'; loading: boolean; }
  | { type: 'account_balance_updated'; accountId: number; balance: Balance, tokensBalance: Record<string, Balance>; }
  | { type: 'vault_balance_updated'; address: string, balance: Balance; }
  | { type: 'main_address_updated'; address: string; };

export type WalletUpdateCallback = (event: WalletEvent) => void;

export class WalletManager {

  private static readonly GAP_LIMIT = 20;
  private static readonly DEBOUNCE_MS = 1000;

  private addressResolvers: Map<string, () => void>;

  public accounts: Map<number, AccountDTO>;
  public accountsAddressToId: Map<string, number>;

  public receiveAddresses: AddressDTO[];
  public changeAddresses: AddressDTO[];

  public vaults: Map<string, VaultDTO>;

  private pendingUpdates: Map<string, UpdateInfo>;
  private updateTimer?: number;

  private updateCallback?: WalletUpdateCallback;

  private readonly discoveryService: WalletDiscoveryService;

  public constructor(
    private readonly keyManager: KeyManager,
    private readonly kvStore: KVStore,
    private readonly walletDb: WalletDB,
    private readonly rostrumService: RostrumService,
    private readonly assetService: AssetService,
    private readonly transactionService: TransactionService,
    private readonly sessionManager: SessionManager,
  ) {
    this.discoveryService = new WalletDiscoveryService(this.rostrumService, this.keyManager);
    this.accounts = new Map();
    this.accountsAddressToId = new Map();
    this.receiveAddresses = [];
    this.changeAddresses = [];
    this.vaults = new Map();
    this.pendingUpdates = new Map();
    this.addressResolvers = new Map();
  }

  public onUpdate(callback: WalletUpdateCallback): void {
    this.updateCallback = callback;
  }

  private notify(event: WalletEvent): void {
    this.updateCallback?.(event);
  }

  private getAllAddresses(): WalletAddress[] {
    return [...this.receiveAddresses, ...this.changeAddresses, ...this.accounts.values(), ...this.vaults.values()];
  }

  public getMainAddresses(): AddressDTO[] {
    return [...this.receiveAddresses, ...this.changeAddresses];
  }

  public getReceiveAddress(): string {
    return this.receiveAddresses.find(addr => !addr.used)!.address;
  }

  public getChangeAddress(accountId = MAIN_WALLET_ID): string {
    if (accountId != MAIN_WALLET_ID) {
      return this.accounts.get(accountId)!.address;
    }
    return this.changeAddresses.find(addr => !addr.used)!.address;
  }

  public getUsedAddressKeys(accountId = MAIN_WALLET_ID): AddressKey[] {
    if (accountId != MAIN_WALLET_ID) {
      const account = this.accounts.get(accountId)!;
      return [{
        address: account.address,
        keyPath: keyPathToString(AccountType.DAPP, 0, account.id),
        balance: account.balance,
        tokensBalance: account.tokensBalance
      }];
    }

    return this.getMainAddresses().filter(addr => addr.used).map(addr => ({
      address: addr.address,
      keyPath: keyPathToString(AccountType.MAIN, addr.space, addr.idx),
      balance: addr.balance,
      tokensBalance: addr.tokensBalance
    }));
  }

  public async reloadRostrum(params: RostrumParams): Promise<void> {
    await this.kvStore.saveRostrumParams(params);
    await this.rostrumService.disconnect(true);
    await this.initRostrum();
    await this.subscribeAddresses(this.getAllAddresses())
  }

  public async initRostrum(): Promise<void> {
    await this.rostrumService.connect();
    await this.rostrumService.subscribeHeaders(tip => this.notify({ type: 'new_tip', tip }));
  }

  public async initialize(): Promise<void> {
    await this.initRostrum();

    const isExist = await this.walletDb.countAddresses();
    if (isExist) {
      await this.loadWallet();
    } else {
      await this.discoverWallet();
    }
  }

  private async discoverWallet(): Promise<void> {
    this.notify({ type: 'discover_wallet', loading: true });

    const rIndexPromise = this.discoveryService.discoverWalletIndex(AccountType.MAIN, KeySpace.RECEIVE);
    const cIndexPromise = this.discoveryService.discoverWalletIndex(AccountType.MAIN, KeySpace.CHANGE);
    const dappIndexPromis = this.discoveryService.discoverWalletIndex(AccountType.DAPP, KeySpace.RECEIVE);
    
    const [rIndex, cIndex, dappIndex] = await Promise.all([rIndexPromise, cIndexPromise, dappIndexPromis]);

    this.receiveAddresses = this.deriveAddresses(KeySpace.RECEIVE, 0, rIndex + WalletManager.GAP_LIMIT);
    this.changeAddresses = this.deriveAddresses(KeySpace.CHANGE, 0, cIndex + WalletManager.GAP_LIMIT);
    this.accounts = this.deriveAccounts(0, dappIndex + 1);
    this.accounts.forEach(entry => this.accountsAddressToId.set(entry.address, entry.id));

    await this.saveAddresses(this.receiveAddresses);
    await this.saveAddresses(this.changeAddresses);
    await this.saveAccounts(this.accounts);

    this.initState();

    await this.initialSync();

    this.notify({ type: 'discover_wallet', loading: false });
    this.notify({ type: 'load_wallet', loading: false });

    await this.rescanVaults();
  }

  private async loadWallet(): Promise<void> {
    this.receiveAddresses = await this.walletDb.getReceiveAddresses();
    this.changeAddresses = await this.walletDb.getChangeAddresses();

    const accs = await this.walletDb.getAccounts();
    this.accounts = new Map(accs.map(a => [a.id, a]));
    this.accountsAddressToId = new Map(accs.map(a => [a.address, a.id]));

    const vaults = await this.walletDb.getVaults();
    this.vaults = new Map(vaults.map(v => [v.address, v]));

    this.initState();

    this.notify({ type: 'load_wallet', loading: false });

    await this.reconnectSync();
  }

  private async initialSync(): Promise<void> {
    await this.subscribeAddresses(this.getAllAddresses(), true);

    const tokensLoadPromises: Promise<void>[] = [];
    tokensLoadPromises.push(this.assetService.syncTokens(MAIN_WALLET_ID, sumTokensBalance(this.getMainAddresses().map(a => a.tokensBalance))));
    for (const account of this.accounts.values()) {
      tokensLoadPromises.push(this.assetService.syncTokens(account.id, account.tokensBalance));
    }
    await Promise.all(tokensLoadPromises);
  }

  public async reconnectSync(): Promise<void> {
    await Promise.all([
      this.sessionManager.reload(this.accounts),
      this.subscribeAddresses(this.getAllAddresses())
    ]);
  }

  private initState(): void {
    const walletAddresses =  this.getMainAddresses();

    this.notify({
      type: 'new_account',
      account: {
        id: MAIN_WALLET_ID,
        name: 'Main Wallet',
        address: this.getReceiveAddress(),
        balance: sumBalance(walletAddresses.map(a => a.balance)),
        tokensBalance: sumTokensBalance(walletAddresses.map(a => a.tokensBalance)),
        sessions: {}
      }
    });

    for (const account of this.accounts.values()) {
      this.notify({
        type: 'new_account',
        account: {
          id: account.id,
          name: account.name,
          address: account.address,
          balance: account.balance,
          tokensBalance: account.tokensBalance,
          sessions: {}
        }
      });
    }

    for (const vault of this.vaults.values()) {
      this.notify({
        type: 'new_vault',
        vault: {
          address: vault.address,
          balance: vault.balance,
          block: vault.block,
          index: vault.idx
        }
      });
    }
  }

  private deriveAddresses(space: KeySpace, startIndex: number, count: number): AddressDTO[] {
    const newAddresses: AddressDTO[] = [];
    for (let i = startIndex; i < startIndex + count; i++) {
      const path = keyPathToString(AccountType.MAIN, space, i);
      const address = this.keyManager.getKey(path).privateKey.toAddress().toString();

      const addressDTO: AddressDTO = {
        address: address,
        space: space,
        idx: i,
        used: false,
        height: 0,
        statusHash: '',
        balance: { confirmed: "0", unconfirmed: "0" },
        tokensBalance: {},
        type: AccountType.MAIN
      };

      newAddresses.push(addressDTO);
    }

    return newAddresses;
  }

  private deriveAccounts(startIndex: number, count: number): Map<number, AccountDTO> {
    const newAccounts = new Map<number, AccountDTO>();
    for (let i = startIndex; i < startIndex + count; i++) {
      const path = keyPathToString(AccountType.DAPP, 0, i);
      const address = this.keyManager.getKey(path).privateKey.toAddress().toString();

      const account: AccountDTO = {
        id: i,
        name: `Account ${i + 1}`,
        address: address,
        height: 0,
        hidden: 0,
        statusHash: '',
        balance: { confirmed: "0", unconfirmed: "0" },
        tokensBalance: {},
        type: AccountType.DAPP
      };

      newAccounts.set(account.id, account);
    }

    return newAccounts;
  }

  private async saveAddresses(addresses: AddressDTO[]): Promise<void> {
    for (const addr of addresses) {
      await this.walletDb.saveAddress(addr);
    }
  }

  private async saveAccounts(accounts: Map<number, AccountDTO>): Promise<void> {
    for (const acc of accounts.values()) {
      await this.walletDb.saveAccount(acc);
    }
  }

  private async subscribeAddresses(addresses: WalletAddress[], isInit = false): Promise<void> {
    const subscriptionPromises = addresses.map(async addr => {
      const result = await this.rostrumService.subscribeAddress(addr.address, this.onSubscribeEvent);
      return { addr, result };
    });

    const initPromises: Promise<void>[] = [];
    const subscriptionResults = await Promise.all(subscriptionPromises);

    for (const { addr, result } of subscriptionResults) {
      if (result && typeof result === 'string' && addr.statusHash != result) {
        if (isInit) {
          const p = new Promise<void>((resolve) => {
            this.addressResolvers.set(addr.address, resolve);
          });
          initPromises.push(p);
        }
        this.registerUpdate({address: addr, result});
      }
    }

    await Promise.all(initPromises);
  }

  private onSubscribeEvent = (data: unknown): void => {
    if (!Array.isArray(data) || data.length < 2) {
      return;
    }

    const [address, hash] = data as [string, string];
    const allAddresses = this.getAllAddresses();
    const addrDTO = allAddresses.find(a => a.address === address);

    if (addrDTO && addrDTO.statusHash !== hash) {
      this.registerUpdate({address: addrDTO, result: hash});
    }
  }

  private registerUpdate(updateInfo: UpdateInfo): void {
    this.pendingUpdates.set(updateInfo.address.address, updateInfo);

    clearTimeout(this.updateTimer);
    this.updateTimer = setTimeout(() => {
      this.processPendingUpdates();
    }, WalletManager.DEBOUNCE_MS);
  }

  private async processPendingUpdates(): Promise<void> {
    if (this.pendingUpdates.size === 0) {
      return;
    }

    try {
      this.notify({ type: 'sync_wallet', loading: true });
      console.log(`Processing ${this.pendingUpdates.size} pending updates...`);

      const updates = Array.from(this.pendingUpdates);
      this.pendingUpdates.clear();

      const accountsHistoryPromises: Promise<AddressHistory>[] = [];
      const vaultsHistoryPromises: Promise<AddressHistory>[] = [];
      const walletHistoryPromises: Promise<AddressHistory>[] = [];
      for (const [, update] of updates) {
        if (this.isAccountAddress(update.address)) {
          const p = this.fetchAndUpdateAccount(update.address, update.result);
          accountsHistoryPromises.push(p);
        } else if (this.isVaultAddress(update.address)) {
          const p = this.fetchAndUpdateVault(update.address, update.result);
          vaultsHistoryPromises.push(p);
        } else {
          const p = this.fetchAndUpdateAddress(update.address, update.result);
          walletHistoryPromises.push(p);
        }
      }
      
      const accountsTxs = await Promise.all(accountsHistoryPromises);
      await Promise.all(vaultsHistoryPromises);
      const walletTxs = await Promise.all(walletHistoryPromises);

      if (this.addressResolvers.size > 0) {
        for (const [address,] of updates) {
          const resolver = this.addressResolvers.get(address);
          if (resolver) {
            resolver();
            this.addressResolvers.delete(address);
          }
        }
      }

      await Promise.all([this.postProcessWalletUpdate(walletTxs), this.postProcessAccountsUpdate(accountsTxs)]);
    } catch (e) {
      console.error('Error processing pending updates:', e);
    } finally {
      this.notify({ type: 'sync_wallet', loading: false });
    }
  }

  private async postProcessWalletUpdate(history: AddressHistory[]): Promise<void> {
    if (history.length === 0) {
      return;
    }

    const txHistoryMap = new Map<string, ITXHistory>();
    for (const { txs } of history) {
      for (const tx of txs) {
        txHistoryMap.set(tx.tx_hash, tx);
      }
    }

    await this.checkGapLimit(KeySpace.RECEIVE);
    await this.checkGapLimit(KeySpace.CHANGE);

    const walletAddresses =  this.getMainAddresses();
    const walletBalance = sumBalance(walletAddresses.map(a => a.balance));
    const walletTokenBalances = sumTokensBalance(walletAddresses.map(a => a.tokensBalance));

    this.notify({ type: 'account_balance_updated', accountId: MAIN_WALLET_ID, balance: walletBalance, tokensBalance: walletTokenBalances });
    this.notify({ type: 'main_address_updated', address: this.getReceiveAddress() });

    await this.postProcessActivity(MAIN_WALLET_ID, Array.from(txHistoryMap.values()), walletAddresses.map(a => a.address), walletTokenBalances);
  }

  private async postProcessAccountsUpdate(history: AddressHistory[]): Promise<void> {
    if (history.length === 0) {
      return;
    }

    const activityPromises: Promise<void>[] = [];
    for (const { address, txs } of history) {
      const account = address as AccountDTO;
      
      this.notify({ type: 'account_balance_updated', accountId: account.id, balance: account.balance, tokensBalance: account.tokensBalance });

      const p = this.postProcessActivity(account.id, txs, [account.address], account.tokensBalance);
      activityPromises.push(p);
    }
    await Promise.all(activityPromises);
  }

  private async postProcessActivity(accountId: number, txHistory: ITXHistory[], addresses: string[], tokensBalance: Record<string, Balance>): Promise<void> {
    const postPromises: Promise<void>[] = [];
    const nftPromise = this.assetService.syncNfts(accountId, tokensBalance);
    for (const tx of txHistory) {
      const p = this.transactionService.classifyAndSaveTransaction(accountId, tx.tx_hash, addresses);
      postPromises.push(p);
    }
    postPromises.push(nftPromise);
    await Promise.all(postPromises);
  }

  private async fetchAndUpdateAddress(addrDTO: AddressDTO, hash: string): Promise<AddressHistory> {
    addrDTO.statusHash = hash;
    addrDTO.balance = await this.rostrumService.getBalance(addrDTO.address);
    addrDTO.tokensBalance = await this.rostrumService.getTokensBalance(addrDTO.address);

    const history = await this.transactionService.fetchTransactionsHistory(addrDTO.address, addrDTO.height);
    addrDTO.height = history.lastHeight;
    if (!addrDTO.used) {
      addrDTO.used = history.txs.length > 0;
    }

    await this.walletDb.saveAddress(addrDTO);

    return { address: addrDTO, txs: history.txs };
  }

  private async fetchAndUpdateAccount(account: AccountDTO, hash: string): Promise<AddressHistory> {
    account.statusHash = hash;
    account.balance = await this.rostrumService.getBalance(account.address);
    account.tokensBalance = await this.rostrumService.getTokensBalance(account.address);

    const history = await this.transactionService.fetchTransactionsHistory(account.address, account.height);
    account.height = history.lastHeight;

    await this.walletDb.saveAccount(account);

    return { address: account, txs: history.txs };
  }

  private async fetchAndUpdateVault(vault: VaultDTO, hash: string): Promise<AddressHistory> {
    vault.statusHash = hash;
    vault.balance = await this.rostrumService.getBalance(vault.address);

    const history = await this.transactionService.fetchTransactionsHistory(vault.address, vault.height);
    vault.height = history.lastHeight;

    await this.walletDb.saveVault(vault);
    this.notify({ type: 'vault_balance_updated', address: vault.address, balance: vault.balance });

    return { address: vault, txs: history.txs };
  }

  private async checkGapLimit(keySpace: KeySpace): Promise<void> {
    const addresses = keySpace === KeySpace.RECEIVE 
      ? this.receiveAddresses 
      : this.changeAddresses;

    const spaceName = keySpace === KeySpace.RECEIVE ? 'Receive' : 'Change';
    const unused = addresses.filter(a => !a.used).length;

    if (unused < WalletManager.GAP_LIMIT) {
      const needed = WalletManager.GAP_LIMIT - unused;
      const lastIndex = addresses[addresses.length - 1].idx;
      
      console.log(`Deriving ${needed} more ${spaceName} addresses...`);
      
      const newAddresses = this.deriveAddresses(keySpace, lastIndex + 1, needed);

      if (keySpace === KeySpace.RECEIVE) {
        this.receiveAddresses.push(...newAddresses);
      } else {
        this.changeAddresses.push(...newAddresses);
      }

      await this.saveAddresses(newAddresses);
      await this.subscribeAddresses(newAddresses);
    }
  }

  public async addNewAccount(id: number, name: string): Promise<void> {
    const path = keyPathToString(AccountType.DAPP, 0, id);
    const address = this.keyManager.getKey(path).privateKey.toAddress().toString();

    const account: AccountDTO = {
      id: id,
      name: name,
      address: address,
      height: 0,
      hidden: 0,
      statusHash: '',
      balance: { confirmed: "0", unconfirmed: "0" },
      tokensBalance: {},
      type: AccountType.DAPP
    };

    this.accounts.set(id, account);
    this.accountsAddressToId.set(address, id);
    await this.walletDb.saveAccount(account);

    this.notify({
      type: 'new_account',
      account: {
        id: account.id,
        name: account.name,
        address: account.address,
        balance: account.balance,
        tokensBalance: account.tokensBalance,
        sessions: {}
      }
    });
    await this.subscribeAddresses([account]);
  }

  public async updateAccountName(id: number, name: string): Promise<void> {
    if (!name) {
      return;
    }

    const account = this.accounts.get(id)!;
    account.name = name;

    await this.walletDb.updateAccountName(id, name);
  }

  public rescanAccount(id: number): void {
    if (id == MAIN_WALLET_ID) {
      const addrs = this.getMainAddresses();
      for (const addr of addrs) {
        addr.height = 0;
        this.registerUpdate({address: addr, result: addr.statusHash});
      }
    } else {
      const account = this.accounts.get(id)!
      account.height = 0;
      this.registerUpdate({address: account, result: account.statusHash});
    }
  }
  
  public getVaultNextIndex(): number {
    if (this.vaults.size == 0) {
      return 0;
    }
    return Math.max(...Array.from(this.vaults.values(), v => v.idx)) + 1;
  }

  public async addNewVault(address: string, block: number, index: number): Promise<void> {
    const vault: VaultDTO = {
      address: address,
      block: block,
      idx: index,
      height: 0,
      statusHash: '',
      balance: { confirmed: "0", unconfirmed: "0" },
      tokensBalance: {},
      type: AccountType.VAULT
    };
    await this.addVault(vault);
  }

  private async addVault(vault: VaultDTO): Promise<void> {
    this.vaults.set(vault.address, vault);
    await this.walletDb.saveVault(vault);
    
    this.notify({
      type: 'new_vault',
      vault: {
        address: vault.address,
        balance: vault.balance,
        block: vault.block,
        index: vault.idx
      }
    });
    await this.subscribeAddresses([vault]);
  }

  public async rescanVaults(): Promise<boolean> {
    let found = false;
    const vaults = await this.discoveryService.discoverVaults(this.getMainAddresses().map(a => a.address));
    const promises: Promise<void>[] = [];
    for (const [address, vault] of vaults) {
      if (!this.vaults.has(address)) {
        found = true;
        promises.push(this.addVault(vault));
      }
    }
    await Promise.all(promises);
    return found;
  }

  private isVaultAddress(address: AccountDTO | AddressDTO | VaultDTO): address is VaultDTO {
    return address.type == AccountType.VAULT;
  }

  private isAccountAddress(address: AccountDTO | AddressDTO | VaultDTO): address is AccountDTO {
    return address.type == AccountType.DAPP;
  }

  private isMainAddress(address: AccountDTO | AddressDTO | VaultDTO): address is AddressDTO {
    return address.type == AccountType.MAIN;
  }
}