import os
import sys
import time
import json
import argparse
from tqdm import tqdm
from web3 import Web3

JML_ADDRESS = "0x91ba8A14D2CC851aBb69212c09f59e06e1e7f0a5"
GMV_ADDRESS = "0x7EF911f8ef130F73D166468c0068753932357B17"

def parse_arguments():
    parser = argparse.ArgumentParser(description="EVM NFT Snapshot Collector for Ghostchain Preclaims")
    parser.add_argument("--rpc", required=True, help="EVM JSON-RPC Node URL")
    parser.add_argument("--delay", type=float, default=0.3, help="Delay between RPC requests in seconds (default: 0.3)")
    return parser.parse_args()


def get_contract_instance(w3):
    abi = [
        {
            "inputs": [],
            "name": "totalSupply",
            "outputs": [{"internalType": "uint256", "name": "", "type": "uint256"}],
            "stateMutability": "view",
            "type": "function"
        },
        {
            "inputs": [{"internalType": "uint256", "name": "_tokenId", "type": "uint256"}],
            "name": "ownerOf",
            "outputs": [{"internalType": "address", "name": "owner", "type": "address"}],
            "stateMutability": "view",
            "type": "function"
        },
        {
            "inputs":[
                {"internalType": "uint256", "name": "", "type":"uint256"},
                {"internalType": "address", "name": "", "type":"address"}
            ],
            "name": "collateralBalances",
            "outputs":[{"internalType": "uint256", "name":"", "type": "uint256"}],
            "stateMutability": "view",
            "type": "function"
        }
    ]
    return w3.eth.contract(address=w3.to_checksum_address(JML_ADDRESS), abi=abi)

def load_existing_snapshot(file_path):
    if os.path.exists(file_path):
        try:
            with open(file_path, "r") as f:
                data = json.load(f)
                if isinstance(data, list) and len(data) > 0:
                    last_token_id = max(item["token_id"] for item in data)
                    print(f"[+] Snapshot exists. Last processed Token ID: {last_token_id}")
                    return data, last_token_id + 1
        except Exception as e:
            print(f"[!] File not fount at {file_path}. Start from scratch.")
    return [], 1

def save_snapshot(file_path, snapshot_data):
    os.makedirs(os.path.dirname(file_path), exist_ok=True)
    with open(file_path, "w") as f:
        json.dump(snapshot_data, f, indent=2)

def main():
    args = parse_arguments()

    w3 = Web3(Web3.HTTPProvider(args.rpc))
    if not w3.is_connected():
        print(f"[-] Could not connect to EVM RPC: {args.rpc}")
        sys.exit(1)

    chain_id = w3.eth.chain_id
    print(f"[+] Successfully connected to [{chain_id}] over {args.rpc}.")

    contract = get_contract_instance(w3)
    total_supply = contract.functions.totalSupply().call()
    print(f"[+] Current totalSupply: {total_supply}")
    time.sleep(args.delay)

    output_file = os.path.join("snapshots", f"{chain_id}.json")
    snapshot, start_token_id = load_existing_snapshot(output_file)

    if start_token_id >= total_supply:
        print(f"[+] All preclaims are already collected.")
        sys.exit(0)

    print(f"[+] Start loop over from Token ID #{start_token_id + 1} to #{total_supply}...")

    try:
        current_index = len(snapshot)
        for token_id in tqdm(range(start_token_id, total_supply), desc=f"Scanning {chain_id}"):
            try:
                owner = contract.functions.ownerOf(token_id).call()
                time.sleep(args.delay)
                shares = contract.functions.collateralBalances(token_id, GMV_ADDRESS).call()
                time.sleep(args.delay)

                snapshot.append({
                    "index": current_index,
                    "token_id": token_id,
                    "address": owner,
                    "shares": shares
                })

                current_index += 1
            except Exception as e:
                print(f"[-] Error occured during on the Token ID {token_id}: {e}")
                break

    except KeyboardInterrupt:
        print("\n[!] Stopped. Saving data...")
    except Exception as fatal_err:
        print(f"\n[-] Critical error: {fatal_err}. Saving data...")
    finally:
        save_snapshot(output_file, snapshot)
        print(f"[+] Saved to file: {output_file}")
        print(f"[+] Chain [{chain_id}] has {len(snapshot)} records")

if __name__ == "__main__":
    main()