import json

from decimal import Decimal
from collections import defaultdict

mainnet_genesis_dir="mainnet-genesis"
testnet_state_file=f"{mainnet_genesis_dir}/testnet-state.json"
mainnet_genesis_file=f"{mainnet_genesis_dir}/config/genesis.json"

with open(testnet_state_file) as f:
  testnet_state = json.load(f)

with open(mainnet_genesis_file) as f:
  mainnet_state = json.load(f)

#------

# Import required module state
mainnet_state["app_state"]["auth"] = testnet_state["app_state"]["auth"]
mainnet_state["app_state"]["bond"] = testnet_state["app_state"]["bond"]
mainnet_state["app_state"]["bank"] = testnet_state["app_state"]["bank"]
mainnet_state["app_state"]["registry"] = testnet_state["app_state"]["registry"]
mainnet_state["app_state"]["slashing"]["params"] = testnet_state["app_state"]["slashing"]["params"]
mainnet_state["consensus"]["params"] = testnet_state["consensus"]["params"]

#------

# Allocate back delegation stakes

# Build map of address -> extra balance to add (as integers)
deltas = {}
for delegation in testnet_state["app_state"]["staking"]["delegations"]:
  addr = delegation["delegator_address"]
  amount = int(Decimal(delegation["shares"]))
  deltas[addr] = deltas.get(addr, 0) + amount

# Now apply the deltas to existing balances
supply_increment = 0
for balance in mainnet_state["app_state"]["bank"]["balances"]:
  addr = balance["address"]
  if addr in deltas:
    for coin in balance["coins"]:
      if coin["denom"] == "alnt":
        coin["amount"] = str(int(coin["amount"]) + deltas[addr])
        supply_increment += deltas[addr]
        break
    del deltas[addr]

# Increase the total supply
for coin in mainnet_state["app_state"]["bank"]["supply"]:
  if coin["denom"] == "alnt":
    coin["amount"] = str(int(coin["amount"]) + supply_increment)

#------

# Remove non-required module accounts

# Addresses to remove
addresses_to_remove = {
  "bonded_tokens_pool",
  "not_bonded_tokens_pool",
  "distribution"
}

# Remove from auth.accounts and get their addresses
removed_addresses = set()
new_accounts = []
for account in mainnet_state["app_state"]["auth"]["accounts"]:
  account_type = account.get("@type", "")
  if "ModuleAccount" in account_type and account.get("name") in addresses_to_remove:
    removed_addresses.add(account["base_account"]["address"])
    continue
  new_accounts.append(account)

mainnet_state["app_state"]["auth"]["accounts"] = new_accounts

# Remove from bank.balances and tally removed amounts
new_balances = []
removed_amounts = defaultdict(int)

for bal in mainnet_state["app_state"]["bank"]["balances"]:
  if bal["address"] in removed_addresses:
    # Skip this account
    for coin in bal["coins"]:
      denom = coin["denom"]
      amount = int(coin["amount"])
      removed_amounts[denom] += amount

    continue

  new_balances.append(bal)

mainnet_state["app_state"]["bank"]["balances"] = new_balances

# Reduce from bank supply
new_supply = []
for coin in mainnet_state["app_state"]["bank"]["supply"]:
  denom = coin["denom"]
  amount = int(coin["amount"])
  amount -= removed_amounts.get(denom, 0)
  new_supply.append({
    "denom": denom,
    "amount": str(amount)
  })

mainnet_state["app_state"]["bank"]["supply"] = new_supply

#------

# Write back modified state
with open(mainnet_genesis_file, "w") as f:
    json.dump(mainnet_state, f, indent=2)