import * as sbv2 from "./src";
import { AggregatorAccount, TransactionObject } from "./src";
import { setupOutputDir } from "./utils";

import * as anchor from "@coral-xyz/anchor";
import { clusterApiUrl, Connection, PublicKey } from "@solana/web3.js";
import { sleep } from "@switchboard-xyz/common";
// import { backOff } from 'exponential-backoff';
import fs from "fs";
import _ from "lodash";
import os from "os";
import path from "path";

const aggregatorMapPath = path.join(
  os.homedir(),
  "devnet-migration",
  sbv2.SB_V2_PID.toBase58(),
  "aggregator_map.csv"
);

async function main() {
  const [oldDirPath, oldFeedDirPath, oldJobDirPath] = setupOutputDir(
    "2TfB33aLaneQb5TNVwyDz3jSZXS6jdW2ARw1Dgf84XCG"
  );
  const [newDirPath, newFeedDirPath, newJobDirPath] = setupOutputDir(
    sbv2.SB_V2_PID.toBase58()
  );

  const devnetConnection = new Connection(
    process.env.SOLANA_DEVNET_RPC ?? clusterApiUrl("devnet")
  );
  console.log(`rpcUrl: ${devnetConnection.rpcEndpoint}`);

  const payer = sbv2.SwitchboardTestContextV2.loadKeypair(
    "~/switchboard_environments_v2/devnet/upgrade_authority/upgrade_authority.json"
  );
  console.log(`payer: ${payer.publicKey.toBase58()}`);

  const newProgram = await sbv2.SwitchboardProgram.load(
    "devnet",
    devnetConnection,
    payer,
    sbv2.SB_V2_PID
  );

  const aggregatorMap = loadAggregatorMap();
  const reverseAggregatorMap = new Map(
    Array.from(aggregatorMap.entries()).map((r) => [r[1], r[0]])
  );

  console.log(
    `Found ${aggregatorMap.size} aggregators that have already been migrated`
  );

  const oldAggregatorKeys = Array.from(aggregatorMap.keys()).map(
    (a) => new PublicKey(a)
  );

  const oldAggregators = await fetchAggregators(
    devnetConnection,
    oldAggregatorKeys,
    new PublicKey("2TfB33aLaneQb5TNVwyDz3jSZXS6jdW2ARw1Dgf84XCG")
  );

  console.log(`Loaded ${oldAggregators.size} old aggregators `);

  const newAggregatorKeys = Array.from(aggregatorMap.values()).map(
    (a) => new PublicKey(a)
  );

  const aggregators = await fetchAggregators(
    devnetConnection,
    newAggregatorKeys,
    sbv2.SB_V2_PID,
    payer.publicKey
  );

  console.log(`Found ${aggregators.size} aggregators to release`);

  const txns: Array<TransactionObject> = [];
  for (const [aggregatorKey, aggregator] of aggregators.entries()) {
    const oldAggregatorKey = reverseAggregatorMap.get(aggregatorKey);
    if (!oldAggregatorKey) {
      console.log(
        `Old aggregator key not found for new aggregator ${aggregatorKey}`
      );
      continue;
    }
    const oldAggregator = oldAggregators.get(oldAggregatorKey);
    if (!oldAggregator) {
      console.log(`Failed to find old aggregator ${oldAggregatorKey}`);
      continue;
    }
    const aggregatorAccount = new AggregatorAccount(
      newProgram,
      new PublicKey(aggregatorKey)
    );
    if (aggregator.resolutionMode.kind === "ModeRoundResolution") {
      txns.push(
        aggregatorAccount.setResolutionModeInstruction(payer.publicKey, {
          authority: payer,
          mode: new sbv2.types.AggregatorResolutionMode.ModeSlidingResolution(),
        })
      );
    }

    txns.push(
      aggregatorAccount.setAuthorityInstruction(payer.publicKey, {
        newAuthority: oldAggregator.authority,
        authority: payer,
      })
    );
  }

  const packedTxns = TransactionObject.pack(txns);

  console.log(`Packed into ${packedTxns.length} txns`);

  const recentBlockhash = await devnetConnection.getLatestBlockhash();

  const signatures = await newProgram.signAndSendAll(
    packedTxns,
    {
      skipConfrimation: true,
      skipPreflight: true,
    },
    recentBlockhash
  );

  for (const signature of signatures) {
    console.log(`https://explorer.solana.com/tx/${signature}?cluster=devnet`);
  }

  await sleep(20000);

  const parsedTxns = await devnetConnection.getParsedTransactions(signatures);

  fs.writeFileSync(
    "./set-authority-txns.json",
    JSON.stringify(packedTxns, undefined, 2)
  );
}

main().catch((error) => {
  console.error(error);
});

function loadAggregatorMap(): Map<string, string> {
  if (!fs.existsSync(aggregatorMapPath)) {
    return new Map();
  }

  const map = new Map();
  const fileString = fs.readFileSync(aggregatorMapPath, "utf-8");
  const fileLines = fileString.split("\n").slice(1);
  fileLines.forEach((r) => {
    const [oldPubkey, newPubkey] = r.split(", ");
    map.set(oldPubkey, newPubkey);
  });

  return map;
}

async function fetchAggregators(
  connection: Connection,
  pubkeys: Array<PublicKey>,
  owner: PublicKey,
  authority?: PublicKey
): Promise<Map<string, sbv2.types.AggregatorAccountData>> {
  const aggregators = new Map<string, sbv2.types.AggregatorAccountData>();

  const batches = _.chunk(pubkeys, 100);
  for (const batch of batches) {
    const accountInfos = await anchor.utils.rpc.getMultipleAccounts(
      connection,
      batch
    );
    for (const account of accountInfos) {
      if (!account) {
        continue;
      }
      if (!account.account.owner.equals(owner)) {
        console.error(
          `Aggregator ${account.publicKey} belongs to wrong program, ${account.account.owner}`
        );
        continue;
      }
      try {
        const aggregator = sbv2.types.AggregatorAccountData.decode(
          account.account.data
        );
        if (!authority || aggregator.authority.equals(authority)) {
          aggregators.set(account.publicKey.toBase58(), aggregator);
        }
      } catch {}
    }
  }

  return aggregators;
}