import os
import json
import getpass
from pathlib import Path
from web3 import Web3
from Crypto.Cipher import AES
from Crypto.Protocol.KDF import PBKDF2
from Crypto.Random import get_random_bytes
from rich import print
from rich.prompt import Prompt
from eth_account import Account

WALLETS_DIR = "wallets"

web3 = None
selected_wallet = None
wallet_data = None

if not os.path.exists(WALLETS_DIR):
    os.makedirs(WALLETS_DIR)


def encrypt_wallet(private_key, password):
    salt = get_random_bytes(16)
    key = PBKDF2(password, salt, dkLen=32)
    cipher = AES.new(key, AES.MODE_GCM)
    ciphertext, tag = cipher.encrypt_and_digest(private_key.encode())

    return {
        'ciphertext': ciphertext.hex(),
        'salt': salt.hex(),
        'nonce': cipher.nonce.hex(),
        'tag': tag.hex()
    }


def decrypt_wallet(data, password):
    salt = bytes.fromhex(data['salt'])
    key = PBKDF2(password, salt, dkLen=32)
    cipher = AES.new(key, AES.MODE_GCM, nonce=bytes.fromhex(data['nonce']))
    return cipher.decrypt_and_verify(bytes.fromhex(data['ciphertext']), bytes.fromhex(data['tag'])).decode()


def connect_to_network():
    global web3
    rpc_url = Prompt.ask("[bold cyan]Введите RPC URL[/]", default="https://ethereum-rpc.publicnode.com")
    web3 = Web3(Web3.HTTPProvider(rpc_url))
    if web3.is_connected():
        print("[green]Успешное подключение к сети Ethereum[/]")
    else:
        print("[red]Ошибка подключения к сети![/]")
        web3 = None


def create_wallet():
    Account.enable_unaudited_hdwallet_features()
    acct = Account.create()
    save_wallet(acct)


def import_wallet():
    print("[bold cyan]Введите приватный ключ: [/]", end='')
    private_key = getpass.getpass(prompt="")
    account = Account.from_key(private_key)
    save_wallet(account)


def save_wallet(account):
    print("[bold cyan]Введите пароль для шифрования кошелька: [/]", end='')
    password = getpass.getpass(prompt="")

    encrypted = encrypt_wallet(account.key.hex(), password)

    filename = os.path.join(WALLETS_DIR, f"{account.address}.json")

    with open(filename, 'w') as f:
        json.dump({"address": account.address, "data": encrypted}, f)
    print(f"[green]Кошелек сохранен: {account.address}[/]")


def select_wallet():
    global selected_wallet, wallet_data
    files = os.listdir(WALLETS_DIR)
    if not files:
        print("[red]Нет доступных кошельков.[/]")
        return

    for i, fname in enumerate(files):
        print(f"[{i}] {Path(fname).stem}")

    try:
        choice = int(Prompt.ask("Выберите номер кошелька"))
    except:
        print("[red]Неверный номер.[/]")
        return

    with open(os.path.join(WALLETS_DIR, files[choice])) as f:
        wallet_data = json.load(f)
        selected_wallet = wallet_data['address']
        print(f"[green]Кошелек выбран: {selected_wallet}[/]")


def get_balance():
    if not selected_wallet:
        print("[red]Кошелек не выбран.[/]")
        return
    balance = web3.eth.get_balance(selected_wallet)
    print(f"Баланс: {web3.from_wei(balance, 'ether')} ETH")


def send_transaction():
    if not selected_wallet:
        print("[red]Кошелек не выбран.[/]")
        return

    print("[bold cyan]Введите пароль от кошелька: [/]", end='')
    password = getpass.getpass(prompt="")

    try:
        private_key = decrypt_wallet(wallet_data['data'], password)
    except:
        print("[red]Неверный пароль![/]")
        return

    to = Prompt.ask("[bold cyan]Введите адрес получателя[/]")
    value = float(Prompt.ask("[bold cyan]Введите сумму в ETH[/]"))

    base_fee = web3.eth.fee_history(1, 'latest')['baseFeePerGas'][-1]
    priority_fee_gwei = float(Prompt.ask("[bold cyan]Введите приоритетную комиссию в GWEI[/]", default="2"))
    max_priority_fee = web3.to_wei(priority_fee_gwei, 'gwei')
    max_fee = base_fee + max_priority_fee * 2

    print(f"[yellow]Максимальная комиссия за транзакцию: {float(max_fee) / 10**9} gwei[/]")

    tx = {
        'chainId': web3.eth.chain_id,
        'nonce': web3.eth.get_transaction_count(selected_wallet),
        'to': to,
        'value': web3.to_wei(value, 'ether'),
        'gas': 21000,
        'maxFeePerGas': max_fee,
        'maxPriorityFeePerGas': max_priority_fee
    }

    try:
        signed = web3.eth.account.sign_transaction(tx, private_key)
        tx_hash = web3.eth.send_raw_transaction(signed.raw_transaction)

        print(f"[green]Транзакция отправлена! TX Hash: {tx_hash.hex()}[/]")
    except Exception as error:
        print(f"[red]{error.message}[/]")


def main():
    print("""
    [bold blue]Меню:[/]
    a. Подключение к сети Ethereum
    c. Создание кошелька
    w. Выбор кошелька
    b. Просмотр баланса
    s. Отправка средств
    i. Импорт существующего кошелька
    q. Выход
            """)

    while True:
        choice = Prompt.ask("Выберите действие").lower()
        if choice == 'a':
            connect_to_network()
        elif choice == 'c':
            create_wallet()
        elif choice == 'w':
            select_wallet()
        elif choice == 'b':
            if web3 and selected_wallet:
                get_balance()
            else:
                print("[red]Необходимо подключиться к сети и выбрать кошелек[/]")
        elif choice == 's':
            if web3 and selected_wallet:
                send_transaction()
            else:
                print("[red]Необходимо подключиться к сети и выбрать кошелек[/]")
        elif choice == 'i':
            import_wallet()
        elif choice == 'q':
            break
        else:
            print("[red]Неверный выбор[/]")


if __name__ == "__main__":
    main()