import json

from decimal import Decimal
from collections import defaultdict
from bech32 import bech32_decode, bech32_encode, convertbits

#------

# Helper methods

def valoper_to_account_address(valoper_addr: str, new_prefix = "laconic") -> str:
  hrp, data = bech32_decode(valoper_addr)
  if hrp is None or data is None:
    raise ValueError("Invalid bech32 address")

  # Convert back from 5-bit to 8-bit
  decoded = convertbits(data, 5, 8, False)
  if decoded is None:
    raise ValueError("Failed to convert bits")

  # Re-encode with new prefix
  converted_data = convertbits(decoded, 8, 5, True)
  return bech32_encode(new_prefix, converted_data)

def get_min_delegation_share(state):
  validators = state["app_state"]["staking"]["validators"]
  delegations = state["app_state"]["staking"]["delegations"]

  shares_list = []

  for val in validators:
    valoper_addr = val["operator_address"]
    orig_delegator_addr = valoper_to_account_address(valoper_addr)

    # Find delegations where the delegator is the original delegator and delegated to their validator
    for delegation in delegations:
      if (delegation["delegator_address"] == orig_delegator_addr and
        delegation["validator_address"] == valoper_addr):
        shares = int(Decimal(delegation["shares"]))
        shares_list.append(shares)

        break

  # Find minimum
  return min(shares_list)

#------

# Read testnet and mainnet states

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"
staking_amount_file=f"output/staking-amount.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

#------

# Find minimum delegation share for common staking amount

min_delegation_share = get_min_delegation_share(testnet_state)
with open(staking_amount_file, "w") as f:
  json.dump({"common_staking_amount": min_delegation_share}, f, indent=2)
print(f"Staking amount written to {staking_amount_file}")

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