import { JsonRpcProvider, Contract } from "ethers";
import { Strategy } from "../src/strategy.schema.js";
import { StrategyCompiler } from "../src/planner/compiler.js";
import { getChainConfig } from "../src/config/chains.js";

export interface SimulationResult {
  success: boolean;
  gasUsed?: bigint;
  error?: string;
  trace?: any;
  stateChanges?: Array<{
    address: string;
    slot: string;
    before: string;
    after: string;
  }>;
}

export async function runForkSimulation(
  strategy: Strategy,
  forkRpc: string,
  blockNumber?: number
): Promise<SimulationResult> {
  const provider = new JsonRpcProvider(forkRpc);
  const chainConfig = getChainConfig(strategy.chain);

  // Create snapshot before simulation
  let snapshotId: string | null = null;
  try {
    snapshotId = await provider.send("evm_snapshot", []);
  } catch (error) {
    // If snapshot not supported, continue without it
    console.warn("Snapshot not supported, continuing without state restore");
  }

  try {
    // Fork at specific block if provided
    if (blockNumber) {
      try {
        await provider.send("anvil_reset", [
          {
            forking: {
              jsonRpcUrl: chainConfig.rpcUrl,
              blockNumber,
            },
          },
        ]);
      } catch (error) {
        // If anvil_reset not available, try hardhat_impersonateAccount or continue
        console.warn("Fork reset not supported, using current state");
      }
    }

    // Compile strategy
    const compiler = new StrategyCompiler(strategy.chain);
    const executorAddr = strategy.executor || process.env.EXECUTOR_ADDR;
    if (!executorAddr) {
      throw new Error("Executor address required for simulation");
    }

    const plan = await compiler.compile(strategy, executorAddr);

    // Execute calls and trace
    const traces: any[] = [];
    const stateChanges: Array<{
      address: string;
      slot: string;
      before: string;
      after: string;
    }> = [];

    for (const call of plan.calls) {
      try {
        // Get state before
        const stateBefore = await getContractState(provider, call.to);

        // Execute call
        const result = await provider.call({
          to: call.to,
          data: call.data,
          value: call.value,
        });

        // Get state after
        const stateAfter = await getContractState(provider, call.to);

        // Record state changes
        for (const slot in stateAfter) {
          if (stateBefore[slot] !== stateAfter[slot]) {
            stateChanges.push({
              address: call.to,
              slot,
              before: stateBefore[slot] || "0x0",
              after: stateAfter[slot],
            });
          }
        }

        traces.push({
          to: call.to,
          data: call.data,
          result,
          success: true,
        });
      } catch (error: any) {
        traces.push({
          to: call.to,
          data: call.data,
          error: error.message,
          success: false,
        });

        // If any call fails, simulation fails
        return {
          success: false,
          error: `Call to ${call.to} failed: ${error.message}`,
          trace: traces,
          stateChanges,
        };
      }
    }

    // Estimate gas
    let gasUsed: bigint | undefined;
    try {
      // Try to get gas estimate from trace
      if (plan.calls.length > 0) {
        const { estimateGasForCalls } = await import("../src/utils/gas.js");
        gasUsed = await estimateGasForCalls(
          provider,
          plan.calls,
          executorAddr
        );
      }
    } catch (error) {
      // Gas estimation failed, continue without it
    }

    return {
      success: true,
      gasUsed,
      trace: traces,
      stateChanges,
    };
  } catch (error: any) {
    return {
      success: false,
      error: error.message,
    };
  } finally {
    // Restore snapshot if available
    if (snapshotId) {
      try {
        await provider.send("evm_revert", [snapshotId]);
      } catch (error) {
        // Ignore revert errors
      }
    }
  }
}

async function getContractState(
  provider: JsonRpcProvider,
  address: string
): Promise<Record<string, string>> {
  // Get storage slots (simplified - in production would get all relevant slots)
  const state: Record<string, string> = {};
  
  // Try to get balance
  try {
    const balance = await provider.getBalance(address);
    state["balance"] = balance.toString();
  } catch {
    // Ignore
  }

  return state;
}