import { encodePacked, keccak256, concatHex, type Address, type Hex } from "viem";

/**
 * A single allowed call, expressed as the tuple the on-chain
 * ManagerWithMerkleVerification hashes into a leaf:
 *
 *   leaf = keccak256(abi.encodePacked(
 *     decoderAndSanitizer, target, valueIsNonZero, selector, packedArgumentAddresses
 *   ))
 *
 * `packedArgumentAddresses` is the tight concatenation of the address arguments the
 * DecoderAndSanitizer returns for this call (e.g. the asset and, where relevant, the
 * recipient — which decoders force to the vault).
 */
export interface LeafTuple {
  decoder: Address;
  target: Address;
  /** true if the manage call sends ETH (value > 0). Defaults to false. */
  valueIsNonZero?: boolean;
  /** 4-byte function selector (0x-prefixed). */
  selector: Hex;
  /** Address arguments the decoder constrains, in order. Packed tight (20 bytes each). */
  argumentAddresses?: Address[];
}

/** Tight-pack address arguments into a single hex blob (matches abi.encodePacked(addr0, addr1, ...)). */
export function packArgumentAddresses(addrs: Address[] = []): Hex {
  if (addrs.length === 0) return "0x";
  return concatHex(addrs);
}

/** Lowercase a 0x address so viem's encodePacked never rejects a non-checksummed (but valid) input.
 *  On-chain the value is 20 raw bytes, so case is irrelevant to the resulting leaf. */
const lc = (a: Address): Address => a.toLowerCase() as Address;

/** Compute a leaf hash exactly as the on-chain verifier does. */
export function leafHash(t: LeafTuple): Hex {
  const packed = packArgumentAddresses((t.argumentAddresses ?? []).map(lc));
  return keccak256(
    encodePacked(
      ["address", "address", "bool", "bytes4", "bytes"],
      [lc(t.decoder), lc(t.target), t.valueIsNonZero ?? false, t.selector, packed],
    ),
  );
}

/**
 * Leaf hash where the packed argument addresses are already the raw bytes the decoder returns
 * (as obtained by eth_call-ing the DecoderAndSanitizer with the call's targetData). This is the
 * production path — the bytes match exactly what ManagerWithMerkleVerification verifies.
 */
export function leafHashWithPacked(t: {
  decoder: Address;
  target: Address;
  valueIsNonZero?: boolean;
  selector: Hex;
  packedArgs: Hex;
}): Hex {
  return keccak256(
    encodePacked(
      ["address", "address", "bool", "bytes4", "bytes"],
      [
        t.decoder.toLowerCase() as Address,
        t.target.toLowerCase() as Address,
        t.valueIsNonZero ?? false,
        t.selector,
        t.packedArgs && t.packedArgs !== "0x" ? t.packedArgs : "0x",
      ],
    ),
  );
}

/** Sorted-pair (commutative) hash of two 32-byte nodes — matches solmate MerkleProofLib.verify. */
export function hashPair(a: Hex, b: Hex): Hex {
  return BigInt(a) <= BigInt(b) ? keccak256(concatHex([a, b])) : keccak256(concatHex([b, a]));
}

export interface MerkleTree {
  root: Hex;
  /** layers[0] = leaves, last layer = [root]. */
  layers: Hex[][];
  leaves: Hex[];
}

/**
 * Build a sorted-pair Merkle tree with odd-node promotion. Self-consistent with
 * MerkleProofLib.verify: any proof from {@link proofFor} reconstructs `root`.
 */
export function buildTree(leaves: Hex[]): MerkleTree {
  if (leaves.length === 0) throw new Error("buildTree: need at least one leaf");
  const layers: Hex[][] = [leaves.slice()];
  let level = leaves.slice();
  while (level.length > 1) {
    const next: Hex[] = [];
    for (let i = 0; i < level.length; i += 2) {
      if (i + 1 < level.length) next.push(hashPair(level[i], level[i + 1]));
      else next.push(level[i]); // odd node promoted unchanged
    }
    layers.push(next);
    level = next;
  }
  return { root: level[0], layers, leaves: leaves.slice() };
}

/** Build a tree directly from tuples. */
export function buildRootFromTuples(tuples: LeafTuple[]): MerkleTree {
  return buildTree(tuples.map(leafHash));
}

/** Merkle proof (bytes32[]) for the leaf at `index` in the tree. */
export function proofFor(tree: MerkleTree, index: number): Hex[] {
  if (index < 0 || index >= tree.leaves.length) throw new Error(`proofFor: index ${index} out of range`);
  const proof: Hex[] = [];
  let idx = index;
  for (let l = 0; l < tree.layers.length - 1; l++) {
    const level = tree.layers[l];
    const isRight = idx % 2 === 1;
    const sibling = isRight ? idx - 1 : idx + 1;
    if (sibling < level.length) proof.push(level[sibling]); // else node was promoted; no element
    idx = Math.floor(idx / 2);
  }
  return proof;
}

/** Proof for a specific leaf hash (looked up by value). */
export function proofForLeaf(tree: MerkleTree, leaf: Hex): Hex[] {
  const idx = tree.leaves.findIndex((l) => l.toLowerCase() === leaf.toLowerCase());
  if (idx === -1) throw new Error("proofForLeaf: leaf not in tree");
  return proofFor(tree, idx);
}

/** Local mirror of MerkleProofLib.verify (sorted-pair). Useful for self-checks/tests. */
export function verify(proof: Hex[], root: Hex, leaf: Hex): boolean {
  let computed = leaf;
  for (const p of proof) computed = hashPair(computed, p);
  return computed.toLowerCase() === root.toLowerCase();
}