import { HDKey } from '@scure/bip32';
import { keccak_256 } from '@noble/hashes/sha3';
import { sha256 } from '@noble/hashes/sha256';
import { bytesToHex } from '@noble/hashes/utils';
import { secp256k1 } from '@noble/curves/secp256k1';
import { createBase58check } from '@scure/base';
import { getAddress } from 'viem';

const b58check = createBase58check(sha256);

const EXTENDED_KEY_VERSIONS: Record<string, { private: number; public: number }> = {
  xpub: { private: 0x0488ade4, public: 0x0488b21e },
  tpub: { private: 0x04358394, public: 0x043587cf },
};

/**
 * Non-custodial tripwire: refuse anything that is not an extended PUBLIC key.
 * This software can only ever receive; it cannot spend.
 */
export function assertPublicOnly(extendedKey: string): void {
  if (extendedKey.slice(0, 4).toLowerCase().includes('prv')) {
    throw new Error(
      'extended key rejected: that is a spending key. This software only accepts public keys (xpub) and can never spend funds.',
    );
  }
}

function parseExtendedPublicKey(extendedKey: string): HDKey {
  assertPublicOnly(extendedKey);
  const prefix = extendedKey.slice(0, 4);
  const versions = EXTENDED_KEY_VERSIONS[prefix];
  if (!versions) {
    throw new Error(`unsupported extended key prefix "${prefix}" — expected xpub or tpub`);
  }
  return HDKey.fromExtendedKey(extendedKey, versions);
}

/** Derive the child public key at the standard external path m/0/index (non-hardened). */
function childPublicKey(extendedKey: string, index: number): Uint8Array {
  if (!Number.isInteger(index) || index < 0 || index >= 2 ** 31) {
    throw new Error(`derivation index out of range: ${index}`);
  }
  const child = parseExtendedPublicKey(extendedKey).deriveChild(0).deriveChild(index);
  if (!child.publicKey) throw new Error('failed to derive public key');
  return child.publicKey;
}

/** keccak256(uncompressed pubkey)[12..] — the 20 address bytes shared by EVM and Tron. */
function addressBytes(compressedPublicKey: Uint8Array): Uint8Array {
  const uncompressed = secp256k1.ProjectivePoint.fromHex(compressedPublicKey).toRawBytes(false).subarray(1);
  return keccak_256(uncompressed).subarray(-20);
}

/** EIP-55 checksummed address for Base at m/0/index. */
export function deriveBaseAddress(xpub: string, index: number): string {
  return getAddress(`0x${bytesToHex(addressBytes(childPublicKey(xpub, index)))}`);
}

/** Base58check Tron address (0x41 prefix) at m/0/index — TIP-32 account-level xpub expected. */
export function deriveTronAddress(xpub: string, index: number): string {
  const payload = new Uint8Array(21);
  payload[0] = 0x41;
  payload.set(addressBytes(childPublicKey(xpub, index)), 1);
  return b58check.encode(payload);
}

export function isValidTronAddress(address: string): boolean {
  try {
    const decoded = b58check.decode(address);
    return decoded.length === 21 && decoded[0] === 0x41;
  } catch {
    return false;
  }
}

export function isValidBaseAddress(address: string): boolean {
  try {
    getAddress(address);
    return true;
  } catch {
    return false;
  }
}