import { Ledger, Test, initializeBindings } from '../../../bindings.js';
import { Types, TypesBigint } from '../../../bindings/mina-transaction/v1/types.js';
import { transactionCommitments } from '../../../mina-signer/src/sign-zkapp-command.js';
import { NetworkId } from '../../../mina-signer/src/types.js';
import { Ml } from '../../ml/conversion.js';
import { PrivateKey, PublicKey } from '../../provable/crypto/signature.js';
import { UInt32, UInt64 } from '../../provable/int.js';
import { Field } from '../../provable/wrapped.js';
import { prettifyStacktrace } from '../../util/errors.js';
import { TupleN } from '../../util/types.js';
import { Actions, Authorization, TokenId, ZkappCommand } from './account-update.js';
import { Account } from './account.js';
import { invalidTransactionError } from './errors.js';
import { type DepthOptions, type TransactionDepthInfo } from './graphql.js';
import {
  Mina,
  defaultNetworkConstants,
  type ActionStates,
  type FeePayerSpec,
} from './mina-instance.js';
import { SimpleLedger } from './transaction-logic/ledger.js';
import {
  defaultNetworkState,
  reportGetAccountError,
  verifyAccountUpdate,
  verifyTransactionLimits,
} from './transaction-validation.js';
import {
  IncludedTransaction,
  PendingTransaction,
  PendingTransactionPromise,
  PendingTransactionStatus,
  RejectedTransaction,
  Transaction,
  WaitForFinalityOptions,
  createRejectedTransaction,
  createTransaction,
  toPendingTransactionPromise,
  toTransactionPromise,
} from './transaction.js';

export { LocalBlockchain, TestPublicKey };

type TestPublicKey = PublicKey & {
  key: PrivateKey;
};
function TestPublicKey(key: PrivateKey): TestPublicKey {
  return Object.assign(PublicKey.fromPrivateKey(key), { key });
}
namespace TestPublicKey {
  export function random<N extends number = 1>(
    count: N = 1 as never
  ): N extends 1 ? TestPublicKey : TupleN<TestPublicKey, N> {
    if (count === 1) return TestPublicKey(PrivateKey.random()) as never;
    return Array.from({ length: count as number }, () =>
      TestPublicKey(PrivateKey.random())
    ) as never;
  }

  export function fromBase58(base58: string): TestPublicKey {
    return TestPublicKey(PrivateKey.fromBase58(base58));
  }
}

export type LocalBlockchain = Awaited<ReturnType<typeof LocalBlockchain>>;

/**
 * A mock Mina blockchain running locally and useful for testing.
 */
async function LocalBlockchain({ proofsEnabled = true, enforceTransactionLimits = true } = {}) {
  await initializeBindings();
  const slotTime = 3 * 60 * 1000;
  const startTime = Date.now();
  const genesisTimestamp = UInt64.from(startTime);
  const ledger = Ledger.create();
  let networkState = defaultNetworkState();

  function addAccount(publicKey: PublicKey, balance: string) {
    try {
      ledger.addAccount(Ml.fromPublicKey(publicKey), balance);
    } catch (error) {
      throw prettifyStacktrace(error);
    }
  }

  let testAccounts = [] as never as TupleN<TestPublicKey, 10>;

  for (let i = 0; i < 10; ++i) {
    let MINA = 10n ** 9n;
    const largeValue = 1000n * MINA;
    const testAccount = TestPublicKey.random();
    addAccount(testAccount, largeValue.toString());
    testAccounts.push(testAccount);
  }

  const events: Record<string, any> = {};
  const actions: Record<string, Record<string, { actions: string[][]; hash: string }[]>> = {};
  const originalProofsEnabled = proofsEnabled;

  return {
    getNetworkId: () => 'devnet' as NetworkId,
    proofsEnabled,
    getNetworkConstants() {
      return {
        ...defaultNetworkConstants,
        genesisTimestamp,
      };
    },
    currentSlot() {
      return UInt32.from(Math.ceil((new Date().valueOf() - startTime) / slotTime));
    },
    hasAccount(publicKey: PublicKey, tokenId: Field = TokenId.default) {
      return !!ledger.getAccount(Ml.fromPublicKey(publicKey), Ml.constFromField(tokenId));
    },
    getAccount(publicKey: PublicKey, tokenId: Field = TokenId.default): Account {
      let accountJson = ledger.getAccount(Ml.fromPublicKey(publicKey), Ml.constFromField(tokenId));
      if (accountJson === undefined) {
        throw new Error(reportGetAccountError(publicKey.toBase58(), TokenId.toBase58(tokenId)));
      }
      return Types.Account.fromJSON(accountJson);
    },
    getNetworkState() {
      return networkState;
    },
    sendTransaction(txn: Transaction<boolean, boolean>): PendingTransactionPromise {
      return toPendingTransactionPromise(async () => {
        let zkappCommandJson = ZkappCommand.toJSON(txn.transaction);
        let commitments = transactionCommitments(
          TypesBigint.ZkappCommand.fromJSON(zkappCommandJson),
          this.getNetworkId()
        );

        if (enforceTransactionLimits) verifyTransactionLimits(txn.transaction);

        // create an ad-hoc ledger to record changes to accounts within the transaction
        let simpleLedger = SimpleLedger.create();

        for (const update of txn.transaction.accountUpdates) {
          let authIsProof = !!update.authorization.proof;
          let kindIsProof = update.body.authorizationKind.isProved.toBoolean();
          // checks and edge case where a proof is expected, but the developer forgot to invoke await tx.prove()
          // this resulted in an assertion OCaml error, which didn't contain any useful information
          if (kindIsProof && !authIsProof) {
            throw Error(
              `The actual authorization does not match the expected authorization kind. Did you forget to invoke \`await tx.prove();\`?`
            );
          }

          let account = simpleLedger.load(update.body);

          // the first time we encounter an account, use it from the persistent ledger
          if (account === undefined) {
            let accountJson = ledger.getAccount(
              Ml.fromPublicKey(update.body.publicKey),
              Ml.constFromField(update.body.tokenId)
            );
            if (accountJson !== undefined) {
              let storedAccount = Account.fromJSON(accountJson);
              simpleLedger.store(storedAccount);
              account = storedAccount;
            }
          }

          // TODO: verify account update even if the account doesn't exist yet, using a default initial account
          if (account !== undefined) {
            let publicInput = update.toPublicInput(txn.transaction);
            await verifyAccountUpdate(
              account,
              update,
              publicInput,
              commitments,
              this.proofsEnabled,
              this.getNetworkId()
            );
            simpleLedger.apply(update);
          }
        }

        let status: PendingTransactionStatus = 'pending';
        const errors: string[] = [];
        try {
          ledger.applyJsonTransaction(
            JSON.stringify(zkappCommandJson),
            defaultNetworkConstants.accountCreationFee.toString(),
            JSON.stringify(networkState)
          );
        } catch (err: any) {
          status = 'rejected';
          try {
            const errorMessages = JSON.parse(err.message);
            const formattedError = invalidTransactionError(txn.transaction, errorMessages, {
              accountCreationFee: defaultNetworkConstants.accountCreationFee.toString(),
            });
            errors.push(formattedError);
          } catch (parseError: any) {
            const fallbackErrorMessage =
              err.message || parseError.message || 'Unknown error occurred';
            errors.push(fallbackErrorMessage);
          }
        }

        // fetches all events from the transaction and stores them
        // events are identified and associated with a publicKey and tokenId
        txn.transaction.accountUpdates.forEach((p, i) => {
          let pJson = zkappCommandJson.accountUpdates[i];
          let addr = pJson.body.publicKey;
          let tokenId = pJson.body.tokenId;
          events[addr] ??= {};
          if (p.body.events.data.length > 0) {
            events[addr][tokenId] ??= [];
            let updatedEvents = p.body.events.data.map((data: Field[]) => {
              return {
                data: data.map((e) => e.toString()),
                transactionInfo: {
                  transactionHash: '',
                  transactionStatus: '',
                  transactionMemo: '',
                },
              };
            });
            events[addr][tokenId].push({
              events: updatedEvents,
              blockHeight: networkState.blockchainLength,
              globalSlot: networkState.globalSlotSinceGenesis,
              // The following fields are fetched from the Mina network. For now, we mock these values out
              // since networkState does not contain these fields.
              blockHash: '',
              parentBlockHash: '',
              chainStatus: '',
            });
          }

          // actions/sequencing events

          // most recent action state
          let storedActions = actions[addr]?.[tokenId];
          let latestActionState_ = storedActions?.[storedActions.length - 1]?.hash;
          // if there exists no hash, this means we initialize our latest hash with the empty state
          let latestActionState =
            latestActionState_ !== undefined
              ? Field(latestActionState_)
              : Actions.emptyActionState();

          actions[addr] ??= {};
          if (p.body.actions.data.length > 0) {
            let newActionState = Actions.updateSequenceState(
              latestActionState,
              p.body.actions.hash
            );
            actions[addr][tokenId] ??= [];
            actions[addr][tokenId].push({
              actions: pJson.body.actions,
              hash: newActionState.toString(),
            });
          }
        });

        let test = await Test();
        const hash = test.transactionHash.hashZkAppCommand(txn.toJSON());
        const pendingTransaction: Omit<PendingTransaction, 'wait' | 'safeWait'> = {
          status,
          errors,
          transaction: txn.transaction,
          setFee: txn.setFee,
          setFeePerSnarkCost: txn.setFeePerSnarkCost,
          hash,
          toJSON: txn.toJSON,
          toPretty: txn.toPretty,
        };

        const wait = async (_options?: {
          maxAttempts?: number;
          interval?: number;
        }): Promise<IncludedTransaction> => {
          const pendingTransaction = await safeWait(_options);
          if (pendingTransaction.status === 'rejected') {
            throw Error(`Transaction failed with errors:\n${pendingTransaction.errors.join('\n')}`);
          }
          return pendingTransaction;
        };

        // Track the block height at which this transaction was included
        const inclusionBlockHeight = Number(networkState.blockchainLength.toBigint());

        const safeWait = async (_options?: {
          maxAttempts?: number;
          interval?: number;
        }): Promise<IncludedTransaction | RejectedTransaction> => {
          if (status === 'rejected') {
            return createRejectedTransaction(pendingTransaction, pendingTransaction.errors);
          }
          return createLocalIncludedTransaction(pendingTransaction, inclusionBlockHeight, () =>
            Number(networkState.blockchainLength.toBigint())
          );
        };

        return {
          ...pendingTransaction,
          wait,
          safeWait,
        };
      });
    },
    transaction(sender: FeePayerSpec, f: () => Promise<void>) {
      return toTransactionPromise(async () => {
        // TODO we run the transaction twice to match the behavior of `Network.transaction`
        let tx = await createTransaction(sender, f, 0, {
          isFinalRunOutsideCircuit: false,
          proofsEnabled: this.proofsEnabled,
          fetchMode: 'test',
        });
        let hasProofs = tx.transaction.accountUpdates.some(Authorization.hasLazyProof);
        return await createTransaction(sender, f, 1, {
          isFinalRunOutsideCircuit: !hasProofs,
          proofsEnabled: this.proofsEnabled,
        });
      });
    },
    applyJsonTransaction(json: string) {
      return ledger.applyJsonTransaction(
        json,
        defaultNetworkConstants.accountCreationFee.toString(),
        JSON.stringify(networkState)
      );
    },
    async fetchEvents(publicKey: PublicKey, tokenId: Field = TokenId.default) {
      // Return events in reverse chronological order (latest events at the beginning)
      const reversedEvents = (
        events?.[publicKey.toBase58()]?.[TokenId.toBase58(tokenId)] ?? []
      ).reverse();
      return reversedEvents;
    },
    async fetchActions(
      publicKey: PublicKey,
      actionStates?: ActionStates,
      tokenId: Field = TokenId.default,
      _from?: number,
      _to?: number
    ) {
      return this.getActions(publicKey, actionStates, tokenId);
    },
    getActions(
      publicKey: PublicKey,
      actionStates?: ActionStates,
      tokenId: Field = TokenId.default
    ): { hash: string; actions: string[][] }[] {
      let currentActions = actions?.[publicKey.toBase58()]?.[TokenId.toBase58(tokenId)] ?? [];
      let { fromActionState, endActionState } = actionStates ?? {};

      let emptyState = Actions.emptyActionState();
      if (endActionState?.equals(emptyState).toBoolean()) return [];

      let start = fromActionState?.equals(emptyState).toBoolean()
        ? undefined
        : fromActionState?.toString();
      let end = endActionState?.toString();

      let startIndex = 0;
      if (start) {
        let i = currentActions.findIndex((e) => e.hash === start);
        if (i === -1) throw Error(`getActions: fromActionState not found.`);
        startIndex = i + 1;
      }
      let endIndex: number | undefined;
      if (end) {
        let i = currentActions.findIndex((e) => e.hash === end);
        if (i === -1) throw Error(`getActions: endActionState not found.`);
        endIndex = i + 1;
      }
      return currentActions.slice(startIndex, endIndex);
    },
    addAccount,
    /**
     * An array of 10 test accounts that have been pre-filled with
     * 30000000000 units of currency.
     */
    testAccounts,
    setGlobalSlot(slot: UInt32 | number) {
      networkState.globalSlotSinceGenesis = UInt32.from(slot);
    },
    incrementGlobalSlot(increment: UInt32 | number) {
      networkState.globalSlotSinceGenesis = networkState.globalSlotSinceGenesis.add(increment);
    },
    setBlockchainLength(height: UInt32) {
      networkState.blockchainLength = height;
    },
    setTotalCurrency(currency: UInt64) {
      networkState.totalCurrency = currency;
    },
    setProofsEnabled(newProofsEnabled: boolean) {
      this.proofsEnabled = newProofsEnabled;
    },
    resetProofsEnabled() {
      this.proofsEnabled = originalProofsEnabled;
    },
  };
}

/**
 * Default finality threshold of 15 blocks provides 99.9% confidence.
 * @see https://docs.minaprotocol.com/mina-protocol/lifecycle-of-a-payment
 */
const DEFAULT_FINALITY_THRESHOLD = 15;

/**
 * Creates an IncludedTransaction with LocalBlockchain-specific depth calculation.
 * Unlike the network version which fetches from GraphQL, this calculates depth
 * based on the current local blockchain state.
 */
function createLocalIncludedTransaction(
  pendingTx: Omit<PendingTransaction, 'wait' | 'safeWait'>,
  inclusionBlockHeight: number,
  getCurrentBlockHeight: () => number
): IncludedTransaction {
  const safeGetDepth = async (options?: DepthOptions): Promise<TransactionDepthInfo | null> => {
    const finalityThreshold = options?.finalityThreshold ?? DEFAULT_FINALITY_THRESHOLD;
    const currentBlockHeight = getCurrentBlockHeight();
    // Depth should never be negative (safeguard against edge cases)
    const depth = Math.max(0, currentBlockHeight - inclusionBlockHeight);

    return {
      depth,
      inclusionBlockHeight,
      currentBlockHeight,
      isFinalized: depth >= finalityThreshold,
      finalityThreshold,
    };
  };

  const getDepth = async (options?: DepthOptions): Promise<TransactionDepthInfo> => {
    const result = await safeGetDepth(options);
    // For LocalBlockchain, this should never be null
    return result!;
  };

  const waitForFinality = async (
    options?: WaitForFinalityOptions
  ): Promise<TransactionDepthInfo> => {
    const finalityThreshold = options?.finalityThreshold ?? DEFAULT_FINALITY_THRESHOLD;
    const interval = options?.interval ?? 60000;
    const maxAttempts = options?.maxAttempts ?? 30;

    for (let attempt = 0; attempt < maxAttempts; attempt++) {
      const info = await safeGetDepth({ finalityThreshold });
      if (info) {
        options?.onProgress?.(info);
        if (info.isFinalized) {
          return info;
        }
      }
      if (attempt < maxAttempts - 1) {
        await new Promise((resolve) => setTimeout(resolve, interval));
      }
    }

    throw Error(
      `Transaction ${pendingTx.hash} did not reach finality after ${maxAttempts} attempts (threshold: ${finalityThreshold} blocks).`
    );
  };

  return {
    status: 'included',
    transaction: pendingTx.transaction,
    toJSON: pendingTx.toJSON,
    toPretty: pendingTx.toPretty,
    hash: pendingTx.hash,
    data: pendingTx.data,
    inclusionBlockHeight,
    getDepth,
    safeGetDepth,
    waitForFinality,
  };
}

// assert type compatibility without preventing LocalBlockchain to return additional properties / methods
LocalBlockchain satisfies (...args: any) => Promise<Mina>;
