import { decode_proof, create_proof } from '@acala-network/chopsticks-executor';
import { Struct, u32, Bytes, Vector, Option, Twox128, blockHeader, u64, Blake2256, compact, Twox64Concat } from '@polkadot-api/substrate-bindings';
import { Binary } from 'polkadot-api';
import { mergeUint8 } from 'polkadot-api/utils';
import { getTxCodec, getExtrinsicDecoder, getConstant, getCallData, unsignedExtrinsic } from '../codecs.js';
import { getSlotDuration, getCurrentSlot, runtimeBlockHeader } from './slot-utils.js';

const textEncoder = new TextEncoder();
const RELAY_CHAIN_SLOT_DURATION_MILLIS = 6000n;
const storagePrefix = (pallet, storage) => Binary.toHex(
  mergeUint8([
    Twox128(textEncoder.encode(pallet)),
    Twox128(textEncoder.encode(storage))
  ])
);
const DMP_QUEUE_HEADS_KEY = storagePrefix("Dmp", "DownwardMessageQueueHeads");
const HRMP_INGRESS_INDEX_KEY = storagePrefix(
  "Hrmp",
  "HrmpIngressChannelsIndex"
);
const HRMP_EGRESS_INDEX_KEY = storagePrefix("Hrmp", "HrmpEgressChannelsIndex");
const HRMP_CHANNELS_KEY = storagePrefix("Hrmp", "HrmpChannels");
const appendParaId = (key, paraId) => Binary.toHex(
  mergeUint8([Binary.fromHex(key), Twox64Concat(u32.enc(paraId))])
);
const hrmpIngressIndexKey = (receiverParaId) => appendParaId(HRMP_INGRESS_INDEX_KEY, receiverParaId);
const hrmpEgressIndexKey = (senderParaId) => appendParaId(HRMP_EGRESS_INDEX_KEY, senderParaId);
const hrmpChannelKey = (sender, receiver) => Binary.toHex(
  mergeUint8([
    Binary.fromHex(HRMP_CHANNELS_KEY),
    Twox64Concat(mergeUint8([u32.enc(sender), u32.enc(receiver)]))
  ])
);
const DEFAULT_CHANNEL = {
  max_capacity: 1e3,
  max_total_size: 102400,
  max_message_size: 102400,
  msg_count: 0,
  total_size: 0,
  mqc_head: void 0
};
const BABE_CURRENT_SLOT_KEY = storagePrefix("Babe", "CurrentSlot");
const PRESERVE_PROOFS = [
  storagePrefix("Babe", "EpochIndex"),
  storagePrefix("Babe", "CurrentBlockRandomness"),
  storagePrefix("Babe", "Randomness"),
  storagePrefix("Babe", "NextRandomness"),
  BABE_CURRENT_SLOT_KEY,
  storagePrefix("Configuration", "ActiveConfig"),
  storagePrefix("Babe", "Authorities")
];
const PARAS_HEADS_PREFIX = Binary.fromHex(storagePrefix("Paras", "Heads"));
const paraHeadKey = (paraId) => {
  const paraIdBytes = u32.enc(paraId);
  return mergeUint8([PARAS_HEADS_PREFIX, Twox64Concat(paraIdBytes)]);
};
const encodeHeadData = (header) => {
  return mergeUint8([compact.enc(header.length), header]);
};
const setValidationDataInherent = async (chain, parentBlock, xcm) => {
  if (!parentBlock.header.digests.some(
    (d) => d.type === "preRuntime" && d.value.engine === "aura"
  ))
    return null;
  const txCodec = getTxCodec(
    parentBlock,
    "ParachainSystem",
    "set_validation_data"
  );
  if (!txCodec) return null;
  const txDec = await getExtrinsicDecoder(parentBlock);
  const prevValidationDataRaw = parentBlock.body.find((raw) => {
    const ext = txDec(raw);
    return ext.call.type === "ParachainSystem" && ext.call.value.type === "set_validation_data";
  });
  const prevValidationDataExt = prevValidationDataRaw && txDec(prevValidationDataRaw);
  const prevValidationData = prevValidationDataExt?.call.value.value.data;
  if (!prevValidationData) {
    throw new Error("TODO no prevValidationData in previous block");
  }
  const paraId = await getConstant(
    parentBlock,
    "ParachainSystem",
    "SelfParaId"
  );
  if (paraId == null) {
    throw new Error("Could not get parachain ID");
  }
  const encodedParentHeader = blockHeader.enc(parentBlock.header);
  const headData = encodeHeadData(encodedParentHeader);
  const storageKey = paraHeadKey(paraId);
  const existingNodes = prevValidationData.relay_chain_state.map((node) => Binary.toHex(node));
  const existingStateRoot = prevValidationData.validation_data.relay_parent_storage_root;
  const decodedProofArray = await decode_proof(
    existingStateRoot,
    existingNodes
  );
  const decodedProof = new Map(decodedProofArray);
  const slotDuration = await getSlotDuration(chain, parentBlock);
  const relaySlotIncreaseRaw = slotDuration / RELAY_CHAIN_SLOT_DURATION_MILLIS;
  const relaySlotIncrease = relaySlotIncreaseRaw > 0n ? relaySlotIncreaseRaw : 1n;
  const relayCurrentSlot = await (async () => {
    const currentSlotHex = decodedProof.get(BABE_CURRENT_SLOT_KEY);
    if (typeof currentSlotHex === "string") {
      return u64.dec(Binary.fromHex(currentSlotHex));
    }
    const currentSlot = await getCurrentSlot(chain, parentBlock);
    return currentSlot * relaySlotIncrease;
  })();
  const nextRelayChainSlot = Binary.toHex(
    u64.enc(relayCurrentSlot + relaySlotIncrease)
  );
  const newEntries = [];
  for (const key of PRESERVE_PROOFS) {
    if (decodedProof.has(key)) {
      const value = key === BABE_CURRENT_SLOT_KEY && nextRelayChainSlot ? nextRelayChainSlot : decodedProof.get(key) ?? null;
      newEntries.push([key, value]);
    }
  }
  newEntries.push([
    Binary.toHex(storageKey),
    Binary.toHex(headData)
  ]);
  const dmpHashKey = appendParaId(DMP_QUEUE_HEADS_KEY, paraId);
  let dmpHash = Binary.fromHex(
    decodedProof.get(dmpHashKey) ?? "0x0000000000000000000000000000000000000000000000000000000000000000"
  );
  for (const { msg, sent_at } of xcm.dmp) {
    dmpHash = Blake2256(
      mqcChain.enc({
        hash: dmpHash,
        sent_at,
        msg_hash: Blake2256(Binary.toOpaque(msg))
      })
    );
  }
  newEntries.push([dmpHashKey, Binary.toHex(dmpHash)]);
  const egressIndexHex = decodedProof.get(hrmpEgressIndexKey(paraId));
  const existingEgressRecipients = egressIndexHex ? vecU32.dec(Binary.fromHex(egressIndexHex)) : [];
  const allEgressRecipients = [
    .../* @__PURE__ */ new Set([...existingEgressRecipients, ...chain.hrmpChannels])
  ].sort((a, b) => a - b);
  if (allEgressRecipients.length > 0) {
    newEntries.push([
      hrmpEgressIndexKey(paraId),
      Binary.toHex(vecU32.enc(allEgressRecipients))
    ]);
    for (const recipientId of allEgressRecipients) {
      const chKey = hrmpChannelKey(paraId, recipientId);
      const existingHex = decodedProof.get(chKey);
      if (existingHex) {
        newEntries.push([chKey, existingHex]);
      } else {
        newEntries.push([
          chKey,
          Binary.toHex(AbridgedHrmpChannelCodec.enc(DEFAULT_CHANNEL))
        ]);
      }
    }
  }
  const nextRelayNumber = Number(prevValidationData.validation_data.relay_parent_number) + Number(relaySlotIncrease);
  const ingressIndexHex = decodedProof.get(hrmpIngressIndexKey(paraId));
  const existingIngressRecipients = ingressIndexHex ? vecU32.dec(Binary.fromHex(ingressIndexHex)) : [];
  const allIngressSenders = [
    .../* @__PURE__ */ new Set([...existingIngressRecipients, ...chain.hrmpChannels])
  ].sort((a, b) => a - b);
  if (allIngressSenders.length > 0)
    newEntries.push([
      hrmpIngressIndexKey(paraId),
      Binary.toHex(vecU32.enc(allIngressSenders))
    ]);
  for (const senderId of allIngressSenders) {
    const chKey = hrmpChannelKey(senderId, paraId);
    const existingHex = decodedProof.get(chKey);
    const ch = existingHex ? AbridgedHrmpChannelCodec.dec(existingHex) : { ...DEFAULT_CHANNEL };
    const hrmpMessages = xcm.hrmp[senderId];
    if (hrmpMessages) {
      let mqcHead = ch.mqc_head ?? new Uint8Array(32);
      let totalSize = ch.total_size;
      for (const data2 of hrmpMessages) {
        mqcHead = Blake2256(
          mqcChain.enc({
            hash: mqcHead,
            sent_at: nextRelayNumber,
            msg_hash: Blake2256(Binary.toOpaque(data2))
          })
        );
        totalSize += data2.length;
      }
      newEntries.push([
        chKey,
        Binary.toHex(
          AbridgedHrmpChannelCodec.enc({
            ...ch,
            msg_count: ch.msg_count + hrmpMessages.length,
            total_size: totalSize,
            mqc_head: mqcHead
          })
        )
      ]);
    } else {
      newEntries.push([chKey, Binary.toHex(AbridgedHrmpChannelCodec.enc(ch))]);
    }
  }
  const [newStateRoot, newNodes] = await create_proof(
    existingNodes,
    newEntries
  );
  const originalDescendants = prevValidationData.relay_parent_descendants ?? [];
  const updatedDescendants = [];
  let lastHeaderHash;
  let nextDescNumber = Number(prevValidationData.validation_data.relay_parent_number) + Number(relaySlotIncrease);
  for (let i = 0; i < originalDescendants.length; i++) {
    const desc = originalDescendants[i];
    const updatedDesc = {
      ...desc,
      // Update state_root for the first descendant
      state_root: i === 0 ? newStateRoot : desc.state_root,
      // Update parentHash to point to the modified previous header
      parent_hash: lastHeaderHash ?? desc.parent_hash,
      // Align descendant numbers with the new relay_parent_number
      number: nextDescNumber
    };
    updatedDescendants.push(updatedDesc);
    nextDescNumber += 1;
    const encoded = runtimeBlockHeader.enc(updatedDesc);
    lastHeaderHash = Binary.toHex(Blake2256(encoded));
  }
  const data = {
    ...prevValidationData,
    validation_data: {
      ...prevValidationData.validation_data,
      relay_parent_number: nextRelayNumber,
      relay_parent_storage_root: newStateRoot
    },
    relay_chain_state: newNodes.map((node) => Binary.fromHex(node)),
    relay_parent_descendants: updatedDescendants
  };
  const inbound_messages_data = {
    downward_messages: {
      full_messages: xcm.dmp,
      hashed_messages: []
    },
    horizontal_messages: {
      full_messages: allIngressSenders.flatMap(
        (senderId) => (xcm.hrmp[senderId] ?? []).map((data2) => [
          senderId,
          {
            sent_at: nextRelayNumber,
            data: data2
          }
        ])
      ),
      hashed_messages: []
    }
  };
  const callData = await getCallData(
    parentBlock,
    "ParachainSystem",
    "set_validation_data",
    {
      data,
      inbound_messages_data
    }
  );
  return unsignedExtrinsic(callData);
};
const mqcChain = Struct({
  hash: Bytes(32),
  sent_at: u32,
  msg_hash: Bytes(32)
});
const vecU32 = Vector(u32);
const AbridgedHrmpChannelCodec = Struct({
  max_capacity: u32,
  max_total_size: u32,
  max_message_size: u32,
  msg_count: u32,
  total_size: u32,
  mqc_head: Option(Bytes(32))
});

export { setValidationDataInherent };
//# sourceMappingURL=set-validation-data.js.map
