import { encodeFunctionData, type Address } from "viem";
import type { LucidlyClient } from "./client.js";
import type { TxResult } from "./types.js";
import { ManagerWithMerkleVerification } from "./artifacts.js";
import { proofFor } from "./merkle.js";
import { LeafNotInRootError, MissingSlippageGuardError } from "./errors.js";
import type { AllowedCall, IntegrationState } from "./integrations.js";

const MAINNET_IDS = new Set<number>([1, 8453, 42161]);

export interface StrategyAction {
  protocol: string;
  action: "approve" | "supply" | "withdraw" | "swap";
  asset: Address | string;
  /** For swaps: the output token. */
  tokenOut?: Address | string;
  amount: bigint;
  /** For swaps: minimum acceptable output (amountOutMinimum). */
  minOut?: bigint;
}

export interface StrategyIntent {
  protocol: string;
  action: "approve" | "supply" | "withdraw" | "swap";
  asset: Address | string;
  tokenOut?: Address | string;
}

export class StrategyModule {
  constructor(private readonly c: LucidlyClient) {}

  private state(vault: Address): IntegrationState {
    const st = this.c.integrations.getState(vault);
    if (!st) {
      throw new Error(
        `No enabled integrations for vault ${vault}. Call integrations.enable(...) first (it sets the Merkle root).`,
      );
    }
    return st;
  }

  private resolve(vault: Address, a: StrategyAction): { call: AllowedCall; amount: bigint; minOut?: bigint } {
    const st = this.state(vault);
    const asset = this.c.resolveToken(a.asset).toLowerCase();
    const tokenOut = a.tokenOut ? this.c.resolveToken(a.tokenOut).toLowerCase() : undefined;
    const call = st.calls.find(
      (c) =>
        c.protocol === a.protocol &&
        c.action === a.action &&
        c.asset.toLowerCase() === asset &&
        (tokenOut === undefined || c.tokenOut?.toLowerCase() === tokenOut),
    );
    if (!call) {
      throw new LeafNotInRootError({ protocol: a.protocol, target: a.tokenOut ? `${asset}->${tokenOut}` : asset, selector: a.action });
    }
    return { call, amount: a.amount, minOut: a.minOut };
  }

  /** Execute a batch of whitelisted actions atomically in one manageVaultWithMerkleVerification tx. */
  async execute(vault: Address, actions: StrategyAction[]): Promise<TxResult> {
    const v = this.c.requireVault("strategy.execute", vault);
    const st = this.state(v);
    const resolved = actions.map((a) => this.resolve(v, a));

    const proofs = resolved.map((r) => proofFor(st.tree, r.call.index));
    const decoders = resolved.map(() => st.decoder);
    const targets = resolved.map((r) => r.call.target);
    const targetData = resolved.map((r) =>
      // decoderFn equals the target's function name (the decoder mirrors the target signature).
      encodeFunctionData({ abi: r.call.abi as any, functionName: r.call.decoderFn, args: r.call.buildExecArgs(r.amount, r.minOut) as any }),
    );
    const values = resolved.map(() => 0n);

    const hash = await this.c.send({
      address: st.manager,
      abi: ManagerWithMerkleVerification.abi,
      functionName: "manageVaultWithMerkleVerification",
      args: [proofs, decoders, targets, targetData, values],
    });
    return this.c.wrapTx(hash);
  }

  /** Check (the agent safety guard) whether a matching leaf exists for an intent. */
  hasLeaf(vault: Address, intent: StrategyIntent): boolean {
    const st = this.c.integrations.getState(vault);
    if (!st) return false;
    const asset = this.c.resolveToken(intent.asset).toLowerCase();
    const tokenOut = intent.tokenOut ? this.c.resolveToken(intent.tokenOut).toLowerCase() : undefined;
    return st.calls.some(
      (c) =>
        c.protocol === intent.protocol &&
        c.action === intent.action &&
        c.asset.toLowerCase() === asset &&
        (tokenOut === undefined || c.tokenOut?.toLowerCase() === tokenOut),
    );
  }

  /**
   * Swap `amount` of `from` into `to` via Uniswap V3 (batches approve + exactInputSingle).
   * On mainnet a slippage guard is required: pass `minOut`, or `allowZeroMinOut: true` to opt out.
   */
  async swap(
    vault: Address,
    params: { from: Address | string; to: Address | string; amount: bigint; minOut?: bigint; allowZeroMinOut?: boolean },
  ): Promise<TxResult> {
    const mainnet = MAINNET_IDS.has(this.c.chainConfig.id);
    if (mainnet && params.minOut === undefined && !params.allowZeroMinOut) {
      throw new MissingSlippageGuardError("strategy.swap");
    }
    return this.execute(vault, [
      { protocol: "UniswapV3", action: "approve", asset: params.from, amount: params.amount },
      { protocol: "UniswapV3", action: "swap", asset: params.from, tokenOut: params.to, amount: params.amount, minOut: params.minOut ?? 0n },
    ]);
  }

  readonly aave = {
    /** Supply `amount` of `asset` into Aave (batches approve + supply in one tx). */
    supply: (vault: Address, params: { asset: Address | string; amount: bigint }): Promise<TxResult> =>
      this.execute(vault, [
        { protocol: "AaveV3", action: "approve", asset: params.asset, amount: params.amount },
        { protocol: "AaveV3", action: "supply", asset: params.asset, amount: params.amount },
      ]),

    /** Withdraw `amount` of `asset` from Aave back to the vault. */
    withdraw: (vault: Address, params: { asset: Address | string; amount: bigint }): Promise<TxResult> =>
      this.execute(vault, [{ protocol: "AaveV3", action: "withdraw", asset: params.asset, amount: params.amount }]),
  };
}