import { toFunctionSelector, type Address, type Hex } from "viem";
import type { LucidlyClient } from "./client.js";
import { LucidlyDecoderAndSanitizer } from "./decoderArtifacts.js";
import { ERC20_ABI, AAVE_V3_POOL_ABI, UNISWAP_V3_ROUTER_ABI } from "./protocolAbis.js";
import { buildTree, leafHashWithPacked, type MerkleTree } from "./merkle.js";
import { regenerateSyUsdBaseline, syUsdLeafDigests } from "./syusd.js";
import { IntegrationUnavailableError } from "./errors.js";

/** A single call whitelisted in the vault's Merkle root. */
export interface AllowedCall {
  key: string;
  protocol: string;
  action: "approve" | "supply" | "withdraw" | "swap";
  asset: Address;
  /** For swaps: the output token. */
  tokenOut?: Address;
  target: Address;
  selector: Hex;
  abi: readonly unknown[];
  decoderFn: string;
  leafArgs: readonly unknown[];
  buildExecArgs: (amount: bigint, minOut?: bigint) => readonly unknown[];
  leaf: Hex;
  packedArgs: Hex;
  index: number;
}

export interface IntegrationState {
  decoder: Address;
  /** Full syUSD decoder the baseline is regenerated against (vault.deploy repointed). */
  syUsdDecoder?: Address;
  strategist: Address;
  manager: Address;
  vault: Address;
  accountant: Address;
  protocols: string[];
  assets: Address[];
  pairs: [Address, Address][];
  fee: number;
  keepSyUsdBaseline: boolean;
  tree: MerkleTree;
  calls: AllowedCall[];
}

/** Minimal, JSON-serializable descriptor sufficient to rebuild the full in-memory state. */
export interface IntegrationDescriptor {
  vault: Address;
  decoder: Address;
  syUsdDecoder?: Address;
  strategist: Address;
  manager: Address;
  accountant: Address;
  protocols: string[];
  assets: Address[];
  pairs: [Address, Address][];
  fee: number;
  keepSyUsdBaseline: boolean;
}

export interface EnableOptions {
  strategist?: Address;
  decoder?: Address;
  /** Assets to whitelist for asset-keyed protocols (AaveV3). Defaults to ["USDC"]. */
  assets?: (Address | string)[];
  /** Token pairs [from, to] to whitelist for swap protocols (UniswapV3). Defaults to [[USDC, WETH]]. */
  pairs?: [Address | string, Address | string][];
  /** Uniswap V3 fee tier used when building swap leaves (the decoder ignores fee, so any tier works). */
  fee?: number;
  /** Keep the syUSD baseline leaves in the root and append the new protocol leaves. Defaults to true. */
  keepSyUsdBaseline?: boolean;
}

export interface EnableResult {
  decoder: Address;
  root: Hex;
  calls: AllowedCall[];
  descriptor: IntegrationDescriptor;
}

type Spec = Omit<AllowedCall, "leaf" | "packedArgs" | "index">;

export class IntegrationsModule {
  private state = new Map<string, IntegrationState>();

  constructor(private readonly c: LucidlyClient) {}

  getState(vault: Address): IntegrationState | undefined {
    return this.state.get(vault.toLowerCase());
  }

  async deployDecoder(vault?: Address): Promise<Address> {
    const v = this.c.requireVault("integrations.deployDecoder", vault);
    // Composite decoder needs the Uniswap V3 NonfungiblePositionManager for LP sanitization.
    const nfpm = this.c.chainConfig.contracts.uniswapV3PositionManager ?? ("0x0000000000000000000000000000000000000000" as Address);
    return this.c.deployContract({
      abi: LucidlyDecoderAndSanitizer.abi,
      bytecode: LucidlyDecoderAndSanitizer.bytecode,
      args: [v, nfpm],
    });
  }

  /** Build the call specs (no leaves yet) for the requested protocols/assets. */
  private buildSpecs(
    protocols: string[],
    assets: Address[],
    vault: Address,
    pairs: [Address, Address][],
    fee: number,
  ): Spec[] {
    const specs: Spec[] = [];
    for (const protocol of protocols) {
      if (protocol === "UniswapV3") {
        const router = this.c.chainConfig.contracts.uniswapV3Router;
        if (!router) throw new IntegrationUnavailableError("UniswapV3", this.c.chainConfig.name);
        for (const [tokenIn, tokenOut] of pairs) {
          // approve the router to spend tokenIn
          specs.push({
            key: `UniswapV3:approve:${tokenIn.toLowerCase()}`, protocol, action: "approve", asset: tokenIn,
            target: tokenIn, selector: toFunctionSelector("approve(address,uint256)"), abi: ERC20_ABI,
            decoderFn: "approve", leafArgs: [router, 0n], buildExecArgs: (a) => [router, a],
          });
          // exactInputSingle swap tokenIn -> tokenOut on SwapRouter02; recipient forced to the vault
          const mkParams = (amountIn: bigint, minOut: bigint) => ({
            tokenIn, tokenOut, fee, recipient: vault, amountIn, amountOutMinimum: minOut, sqrtPriceLimitX96: 0n,
          });
          specs.push({
            key: `UniswapV3:swap:${tokenIn.toLowerCase()}:${tokenOut.toLowerCase()}`, protocol, action: "swap", asset: tokenIn, tokenOut,
            target: router, selector: toFunctionSelector("exactInputSingle((address,address,uint24,address,uint256,uint256,uint160))"), abi: UNISWAP_V3_ROUTER_ABI,
            decoderFn: "exactInputSingle", leafArgs: [mkParams(0n, 0n)],
            buildExecArgs: (a, minOut) => [mkParams(a, minOut ?? 0n)],
          });
        }
      } else if (protocol === "AaveV3") {
        const pool = this.c.chainConfig.contracts.aaveV3Pool;
        if (!pool) throw new IntegrationUnavailableError("AaveV3", this.c.chainConfig.name);
        for (const asset of assets) {
          specs.push({
            key: `AaveV3:approve:${asset.toLowerCase()}`, protocol, action: "approve", asset,
            target: asset, selector: toFunctionSelector("approve(address,uint256)"), abi: ERC20_ABI,
            decoderFn: "approve", leafArgs: [pool, 0n], buildExecArgs: (a) => [pool, a],
          });
          specs.push({
            key: `AaveV3:supply:${asset.toLowerCase()}`, protocol, action: "supply", asset,
            target: pool, selector: toFunctionSelector("supply(address,uint256,address,uint16)"), abi: AAVE_V3_POOL_ABI,
            decoderFn: "supply", leafArgs: [asset, 0n, vault, 0], buildExecArgs: (a) => [asset, a, vault, 0],
          });
          specs.push({
            key: `AaveV3:withdraw:${asset.toLowerCase()}`, protocol, action: "withdraw", asset,
            target: pool, selector: toFunctionSelector("withdraw(address,uint256,address)"), abi: AAVE_V3_POOL_ABI,
            decoderFn: "withdraw", leafArgs: [asset, 0n, vault], buildExecArgs: (a) => [asset, a, vault],
          });
        }
      } else {
        throw new IntegrationUnavailableError(protocol, this.c.chainConfig.name);
      }
    }
    return specs;
  }

  /** Derive packed args + leaf for each spec by eth_call-ing the decoder. */
  private async deriveCalls(decoder: Address, specs: Spec[]): Promise<AllowedCall[]> {
    const calls: AllowedCall[] = [];
    for (let i = 0; i < specs.length; i++) {
      const s = specs[i];
      const packedArgs = await this.c.read<Hex>({
        address: decoder,
        abi: LucidlyDecoderAndSanitizer.abi,
        functionName: s.decoderFn,
        args: s.leafArgs,
      });
      const leaf = leafHashWithPacked({ decoder, target: s.target, valueIsNonZero: false, selector: s.selector, packedArgs });
      calls.push({ ...s, packedArgs, leaf, index: i });
    }
    return calls;
  }

  /**
   * Build the manage tree from `baselineDigests` (the kept syUSD whitelist, repointed to this
   * vault's decoder) followed by the new protocol leaves. Reindexes each call to its position in
   * the combined tree so proofFor(tree, call.index) is correct.
   */
  private combine(calls: AllowedCall[], baselineDigests: readonly Hex[]): MerkleTree {
    const offset = baselineDigests.length;
    calls.forEach((c, i) => {
      c.index = offset + i;
    });
    return buildTree([...baselineDigests, ...calls.map((c) => c.leaf)]);
  }

  /**
   * syUSD baseline digests to keep. If a full syUSD decoder is known (repointed deploy), regenerate
   * the baseline against it (executable). Otherwise fall back to the published (verbatim) digests.
   */
  private baselineDigestsFor(
    keep: boolean,
    vault: Address,
    syUsdDecoder: Address | undefined,
    accountant: Address,
  ): readonly Hex[] {
    if (!keep) return [];
    if (syUsdDecoder) return regenerateSyUsdBaseline({ vault, decoder: syUsdDecoder, accountant }).digests;
    return syUsdLeafDigests();
  }

  /**
   * Enable protocols on a vault: deploy/register the decoder, build the allowed leaves,
   * and set the Merkle root for the strategist.
   */
  async enable(protocols: string[], opts: EnableOptions = {}, vault?: Address): Promise<EnableResult> {
    const v = this.c.requireVault("integrations.enable", vault);
    const manager = this.c.requireManager("integrations.enable");
    const { account } = this.c.requireSigner("integrations.enable");
    const strategist = opts.strategist ?? account.address;
    const assets = (opts.assets ?? ["USDC"]).map((a) => this.c.resolveToken(a));
    const fee = opts.fee ?? 500;
    const pairs: [Address, Address][] = (opts.pairs ?? [["USDC", "WETH"]]).map(
      ([a, b]) => [this.c.resolveToken(a), this.c.resolveToken(b)] as [Address, Address],
    );
    // Convenience decoder for the new protocol leaves (composite ERC20+Aave+Uniswap).
    const decoder = opts.decoder ?? this.c.activeDeployment?.decoder ?? (await this.deployDecoder(v));
    // Full syUSD decoder for the kept baseline (from a repointed deploy), if any.
    const syUsdDecoder = this.c.activeDeployment?.syUsdDecoder;
    const accountant = this.c.requireAccountant("integrations.enable");

    const keepSyUsdBaseline = opts.keepSyUsdBaseline ?? true;
    const specs = this.buildSpecs(protocols, assets, v, pairs, fee);
    const calls = await this.deriveCalls(decoder, specs);
    const baseline = this.baselineDigestsFor(keepSyUsdBaseline, v, syUsdDecoder, accountant);
    const tree = this.combine(calls, baseline);
    await this.c.roles.setManageRoot(strategist, tree.root, manager);

    const descriptor: IntegrationDescriptor = { vault: v, decoder, syUsdDecoder, strategist, manager, accountant, protocols, assets, pairs, fee, keepSyUsdBaseline };
    this.state.set(v.toLowerCase(), { ...descriptor, tree, calls });
    return { decoder, root: tree.root, calls, descriptor };
  }

  /** Rebuild in-memory state from a saved descriptor (no deploy, no setManageRoot). */
  async rehydrate(d: IntegrationDescriptor): Promise<IntegrationState> {
    const specs = this.buildSpecs(d.protocols, d.assets, d.vault, d.pairs ?? [], d.fee ?? 500);
    const calls = await this.deriveCalls(d.decoder, specs);
    const baseline = this.baselineDigestsFor(d.keepSyUsdBaseline ?? true, d.vault, d.syUsdDecoder, d.accountant);
    const tree = this.combine(calls, baseline);
    const state: IntegrationState = { ...d, tree, calls };
    this.state.set(d.vault.toLowerCase(), state);
    return state;
  }

  /** Extract the JSON-serializable descriptor for a vault's enabled integrations. */
  serialize(vault?: Address): IntegrationDescriptor | undefined {
    const v = this.c.requireVault("integrations.serialize", vault);
    const s = this.state.get(v.toLowerCase());
    if (!s) return undefined;
    return { vault: s.vault, decoder: s.decoder, syUsdDecoder: s.syUsdDecoder, strategist: s.strategist, manager: s.manager, accountant: s.accountant, protocols: s.protocols, assets: s.assets, pairs: s.pairs, fee: s.fee, keepSyUsdBaseline: s.keepSyUsdBaseline };
  }

  leaves(vault?: Address): AllowedCall[] {
    const v = this.c.requireVault("integrations.leaves", vault);
    return this.state.get(v.toLowerCase())?.calls ?? [];
  }
}