import {
  AddressLookupTableAccount,
  PublicKey,
  TransactionInstruction,
  VersionedMessage,
  VersionedTransaction,
} from '@solana/web3.js';
import {
  Account,
  AccountMeta,
  AccountRole,
  Address,
  address,
  CompiledTransactionMessage,
  getCompiledTransactionMessageDecoder,
  Instruction,
  isSome,
  lamports,
  none,
  some,
  Transaction,
} from '@solana/kit';
import { ADDRESS_LOOKUP_TABLE_PROGRAM_ADDRESS, AddressLookupTable } from '@solana-program/address-lookup-table';
import { fromVersionedTransaction } from '@solana/compat';

export function fromLegacyPublicKey(legacy: PublicKey): Address {
  return address(legacy.toString());
}

export function toLegacyPublicKey(address: Address): PublicKey {
  return new PublicKey(address.toString());
}

export function fromLegacyInstructions(...legacy: TransactionInstruction[]): Instruction[] {
  return legacy.map(fromLegacyInstruction);
}

export function fromLegacyInstruction(legacy: TransactionInstruction): Instruction {
  const accounts: AccountMeta[] = legacy.keys.map((key) => ({
    address: fromLegacyPublicKey(key.pubkey),
    role: determineRole(key.isSigner, key.isWritable),
  }));
  return {
    programAddress: fromLegacyPublicKey(legacy.programId),
    accounts,
    data: legacy.data,
  };
}

export function toLegacyInstructions(...instructions: Instruction[]): TransactionInstruction[] {
  return instructions.map(toLegacyInstruction);
}

function toLegacyInstruction(instruction: Instruction): TransactionInstruction {
  const keys = (instruction.accounts || []).map((account) => ({
    pubkey: toLegacyPublicKey(account.address),
    isSigner: account.role === AccountRole.WRITABLE_SIGNER || account.role === AccountRole.READONLY_SIGNER,
    isWritable: account.role === AccountRole.WRITABLE_SIGNER || account.role === AccountRole.WRITABLE,
  }));
  return new TransactionInstruction({
    keys,
    programId: toLegacyPublicKey(instruction.programAddress),
    data: instruction.data ? Buffer.from(instruction.data) : undefined,
  });
}

export function fromLegacyLookupTables(...account: AddressLookupTableAccount[]): Account<AddressLookupTable>[] {
  return account.map(fromLegacyLookupTable);
}

export function fromLegacyLookupTable(lut: AddressLookupTableAccount): Account<AddressLookupTable> {
  return {
    address: fromLegacyPublicKey(lut.key),
    programAddress: ADDRESS_LOOKUP_TABLE_PROGRAM_ADDRESS,
    executable: false,
    lamports: lamports(0n),
    space: 0n,
    data: {
      padding: 0,
      lastExtendedSlotStartIndex: lut.state.lastExtendedSlotStartIndex,
      lastExtendedSlot: BigInt(lut.state.lastExtendedSlot),
      addresses: lut.state.addresses.map(fromLegacyPublicKey),
      authority: lut.state.authority ? some(fromLegacyPublicKey(lut.state.authority)) : none(),
      discriminator: 0,
      deactivationSlot: lut.state.deactivationSlot,
    },
  };
}

export function toLegacyLookupTables(...account: Account<AddressLookupTable>[]): AddressLookupTableAccount[] {
  return account.map(toLegacyLookupTable);
}

function toLegacyLookupTable(account: Account<AddressLookupTable>): AddressLookupTableAccount {
  return new AddressLookupTableAccount({
    key: toLegacyPublicKey(account.address),
    state: {
      lastExtendedSlotStartIndex: account.data.lastExtendedSlotStartIndex,
      lastExtendedSlot: Number(account.data.lastExtendedSlot),
      addresses: account.data.addresses.map(toLegacyPublicKey),
      authority: isSome(account.data.authority)
        ? toLegacyPublicKey(account.data.authority.value as Address)
        : undefined,
      deactivationSlot: account.data.deactivationSlot,
    },
  });
}

export function determineRole(isSigner: boolean, isWritable: boolean): AccountRole {
  if (isSigner && isWritable) return AccountRole.WRITABLE_SIGNER;
  if (isSigner) return AccountRole.READONLY_SIGNER;
  if (isWritable) return AccountRole.WRITABLE;
  return AccountRole.READONLY;
}

export function fromLegacyVersionedTransaction(
  transaction: VersionedTransaction,
  lookupTableAccounts: Account<AddressLookupTable>[],
): Transaction | undefined {
  const hasALTs = lookupTableAccounts.length > 0;

  return hasALTs
    ? fromLegacyVersionedTransactionWithAddressLookupTable(transaction, toLegacyLookupTables(...lookupTableAccounts))
    : fromVersionedTransaction(transaction);
}

export function fromLegacyVersionedTransactionWithAddressLookupTable(
  transaction: VersionedTransaction,
  addressLookupTableAccounts: AddressLookupTableAccount[],
): Transaction {
  // Temporarily override getAccountKeys method
  const originalGetAccountKeys = transaction.message.getAccountKeys.bind(transaction.message);

  try {
    // Override with resolved version
    transaction.message.getAccountKeys = function (_args?: {
      addressLookupTableAccounts?: AddressLookupTableAccount[];
    }) {
      return originalGetAccountKeys({ addressLookupTableAccounts });
    };

    const result = fromVersionedTransaction(transaction);

    // Restore before returning
    transaction.message.getAccountKeys = originalGetAccountKeys;
    return result;
  } catch (error) {
    // Restore on error
    transaction.message.getAccountKeys = originalGetAccountKeys;
    throw error;
  }
}

export function toLegacyVersionedTransaction(
  transaction: Transaction,
  lookupTableAccounts: Account<AddressLookupTable>[],
): VersionedTransaction {
  // Deserialize the message bytes to get the VersionedMessage
  const versionedMessage = VersionedMessage.deserialize(Buffer.from(transaction.messageBytes));

  // Create the versioned transaction from the message
  const versionedTransaction = new VersionedTransaction(versionedMessage);

  const addressLookupTableAccounts = lookupTableAccounts.map(toLegacyLookupTable);

  let staticAccountKeys;
  let numRequiredSignatures;

  if (addressLookupTableAccounts.length > 0) {
    // Get account keys with resolved lookup tables
    const accountKeys = versionedMessage.getAccountKeys({ addressLookupTableAccounts });
    staticAccountKeys = accountKeys.staticAccountKeys;
    numRequiredSignatures = versionedMessage.header.numRequiredSignatures;
  } else {
    // No LUTs, use the simpler version
    const accountKeys = versionedMessage.getAccountKeys();
    staticAccountKeys = accountKeys.staticAccountKeys;
    numRequiredSignatures = versionedMessage.header.numRequiredSignatures;
  }

  // Create signatures array in the correct order (reverse of convertSignatures)
  const signaturesArray: Uint8Array[] = [];

  for (let i = 0; i < numRequiredSignatures; i++) {
    const signerAddress = staticAccountKeys[i].toBase58();
    const signature = transaction.signatures[address(signerAddress)];

    if (signature === null || signature === undefined) {
      // Convert null signatures back to all-zero signatures
      signaturesArray.push(new Uint8Array(64)); // all zeros
    } else {
      // Convert SignatureBytes back to Uint8Array
      signaturesArray.push(new Uint8Array(signature));
    }
  }

  // Validate signature count matches requirements
  if (signaturesArray.length !== numRequiredSignatures) {
    throw new Error(`Signature count mismatch: expected ${numRequiredSignatures}, got ${signaturesArray.length}`);
  }

  // Set the signatures on the versioned transaction
  versionedTransaction.signatures = signaturesArray;

  return versionedTransaction;
}

export function extractLutAddressesFromTransaction(transaction: Transaction): Address[] {
  try {
    const dMessage: CompiledTransactionMessage = getCompiledTransactionMessageDecoder().decode(
      transaction.messageBytes,
    );

    if (dMessage.version === 'legacy') {
      return [];
    }

    return dMessage.addressTableLookups?.map((lut) => lut.lookupTableAddress) ?? [];
  } catch (error) {
    console.warn('Failed to extract LUT addresses from transaction:', error);
    return [];
  }
}