import type { Hex, Address } from "viem";
import { buildTree, proofFor, leafHashWithPacked, type MerkleTree } from "./merkle.js";
import {
  SYUSD_ROOT,
  SYUSD_METADATA,
  SYUSD_LEAF_DIGESTS,
  SYUSD_LEAVES,
  type SyUsdLeaf,
} from "./syusdBaseline.js";

export { SYUSD_ROOT, SYUSD_METADATA, SYUSD_LEAVES, type SyUsdLeaf };

let _tree: MerkleTree | null = null;

/** The syUSD baseline Merkle tree (memoized). Its root equals the published SYUSD_ROOT. */
export function syUsdTree(): MerkleTree {
  if (!_tree) _tree = buildTree([...SYUSD_LEAF_DIGESTS]);
  return _tree;
}

/** Ordered syUSD leaf digests (real + padding), used to extend a vault's tree while keeping syUSD. */
export function syUsdLeafDigests(): readonly Hex[] {
  return SYUSD_LEAF_DIGESTS;
}

/** Merkle proof for the syUSD leaf at `index` against SYUSD_ROOT. */
export function syUsdProofFor(index: number): Hex[] {
  return proofFor(syUsdTree(), index);
}

/** A syUSD leaf re-pointed to a target vault (own decoder + recipients). */
export interface RepointedLeaf extends SyUsdLeaf {
  digest: Hex;
}

export interface RegeneratedSyUsd {
  /** Manage root over the regenerated leaves (differs from SYUSD_ROOT). */
  root: Hex;
  /** Full 1024 ordered digests (real leaves regenerated, padding unchanged). */
  digests: Hex[];
  /** The 227 regenerated real leaves. */
  leaves: RepointedLeaf[];
  /** Number of address substitutions applied (decoder field + arg replacements). */
  substitutions: number;
}

/** Replace any 20-byte address chunk in tightly-packed args using the substitution map. */
function substitutePacked(packed: Hex, map: Map<string, string>): { packed: Hex; hits: number } {
  if (!packed || packed === "0x") return { packed: "0x", hits: 0 };
  const hex = packed.slice(2);
  let out = "0x";
  let hits = 0;
  for (let i = 0; i < hex.length; i += 40) {
    const chunk = hex.slice(i, i + 40);
    const repl = map.get(("0x" + chunk).toLowerCase());
    if (repl) {
      out += repl.slice(2);
      hits++;
    } else {
      out += chunk;
    }
  }
  return { packed: out as Hex, hits };
}

/**
 * Regenerate the syUSD strategy set for a specific vault: every leaf's decoder field becomes the
 * target decoder, and syUSD's vault/accountant addresses in the arguments are re-pointed to the
 * target vault/accountant (recipient/onBehalfOf -> target vault). Produces a new root + valid proofs.
 */
export function regenerateSyUsdBaseline(target: {
  vault: Address;
  decoder: Address;
  accountant?: Address;
}): RegeneratedSyUsd {
  const map = new Map<string, string>();
  map.set(SYUSD_METADATA.vault.toLowerCase(), target.vault.toLowerCase());
  map.set(SYUSD_METADATA.decoder.toLowerCase(), target.decoder.toLowerCase());
  if (target.accountant) map.set(SYUSD_METADATA.accountant.toLowerCase(), target.accountant.toLowerCase());

  const digests = [...SYUSD_LEAF_DIGESTS];
  const leaves: RepointedLeaf[] = [];
  let substitutions = 0;
  for (const l of SYUSD_LEAVES) {
    const { packed, hits } = substitutePacked(l.packedArgs, map);
    substitutions += hits + 1; // +1: the decoder field is always re-pointed
    const digest = leafHashWithPacked({
      decoder: target.decoder,
      target: l.target,
      valueIsNonZero: l.canSendValue,
      selector: l.selector,
      packedArgs: packed,
    });
    digests[l.index] = digest;
    leaves.push({ ...l, decoder: target.decoder, packedArgs: packed, digest });
  }
  const tree = buildTree(digests);
  return { root: tree.root, digests, leaves, substitutions };
}

/** Find a syUSD whitelisted leaf by description substring, selector, or target. */
export function findSyUsdLeaf(q: { description?: string; selector?: string; target?: string }): SyUsdLeaf | undefined {
  return SYUSD_LEAVES.find(
    (l) =>
      (q.description === undefined || l.description.toLowerCase().includes(q.description.toLowerCase())) &&
      (q.selector === undefined || l.selector.toLowerCase() === q.selector.toLowerCase()) &&
      (q.target === undefined || l.target.toLowerCase() === q.target.toLowerCase()),
  );
}