/*
 * This file is part of the Monero P2Pool <https://github.com/SChernykh/p2pool>
 * Copyright (c) 2021-2025 SChernykh <https://github.com/SChernykh>
 * Portions Copyright (c) 2012-2013 The Cryptonote developers
 * Portions Copyright (c) 2014-2021 The Monero Project
 * Portions Copyright (c) 2021 XMRig <https://github.com/xmrig>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, version 3.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
 */

#include "common.h"
#include "wallet.h"
#include "keccak.h"
#include "carrot_crypto.h"
#include "crypto.h"
#include <ios>

extern "C" {
#include "crypto-ops.h"
}

#include "fcmp_pp_crypto.h"

LOG_CATEGORY(Wallet)

namespace {

// Allow only regular addresses (no integrated addresses, no subaddresses)
// Values taken from cryptonote_config.h (CRYPTONOTE_PUBLIC_ADDRESS_BASE58_PREFIX)

constexpr uint64_t valid_prefixes[] = { 0x180c96, 0x254c96, 0x24cc96 };  // SC1, SC1T, SC1S (mainnet, testnet, stagenet)
constexpr uint64_t valid_prefixes_subaddress[] = { 0x314c96, 0x3c54c96, 0x384cc96 };  // SC1s, SC1Ts, SC1Ss

constexpr std::array<int, 9> block_sizes{ 0, 2, 3, 5, 6, 7, 9, 10, 11 };
constexpr int num_full_blocks = p2pool::Wallet::ADDRESS_LENGTH / block_sizes.back();
constexpr int last_block_size = p2pool::Wallet::ADDRESS_LENGTH % block_sizes.back();

constexpr int block_sizes_lookup[11] = { 0, -1, 1, 2, -1, 3, 4, 5, -1, 6, 7 };

constexpr char alphabet[] = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz";
constexpr size_t alphabet_size = sizeof(alphabet) - 1;

static_assert(alphabet_size == 58, "Check alphabet");

struct ReverseAlphabet
{
	int8_t data[256];
	int num_symbols;

	static constexpr ReverseAlphabet init()
	{
		ReverseAlphabet result = {};

		for (int i = 0; i < 256; ++i) {
			result.data[i] = -1;
		}

		result.num_symbols = 0;
		for (size_t i = 0; i < alphabet_size; ++i) {
			if (result.data[static_cast<uint8_t>(alphabet[i])] < 0) {
				result.data[static_cast<uint8_t>(alphabet[i])] = static_cast<int8_t>(i);
				++result.num_symbols;
			}
		}

		return result;
	}
};

constexpr ReverseAlphabet rev_alphabet = ReverseAlphabet::init();

static_assert(rev_alphabet.num_symbols == 58, "Check alphabet");

} // namespace

namespace p2pool {

Wallet::Wallet(const char* address) : m_prefix(0), m_checksum(0), m_type(NetworkType::Invalid), m_subaddress(false)
{
	if (!decode(address) && address) {
		LOGWARN(1, address << " failed to decode");
	}
}

Wallet::Wallet(const Wallet& w)
{
	operator=(w);
}

Wallet& Wallet::operator=(const Wallet& w)
{
	if (this == &w) {
		return *this;
	}

	m_prefix = w.m_prefix;
	m_spendPublicKey = w.m_spendPublicKey;
	m_viewPublicKey = w.m_viewPublicKey;
	m_checksum = w.m_checksum;
	m_type = w.m_type;
	m_subaddress = w.m_subaddress;

	return *this;
}

bool Wallet::decode(const char* address)
{
        m_type = NetworkType::Invalid;

        if (!address) {
                return false;
        }

        const size_t addr_len = strlen(address);
        if (addr_len < 94 || addr_len > 140) {
                return false;
        }

        // Calculate based on actual address length
        const int actual_num_full_blocks = addr_len / block_sizes.back();
        const int actual_last_block_size = addr_len % block_sizes.back();
        const int actual_last_block_size_index = block_sizes_lookup[actual_last_block_size];

        if (actual_last_block_size_index < 0) {
                return false;
        }

        uint8_t data[73] = {};
        int data_index = 0;

        const char* addr_ptr = address;  // Use separate pointer for iteration

        for (int i = 0; i <= actual_num_full_blocks; ++i) {
                uint64_t num = 0;
                uint64_t order = 1;

                for (int j = ((i < actual_num_full_blocks) ? block_sizes.back() : actual_last_block_size) - 1; j >= 0; --j) {
                        const int8_t digit = rev_alphabet.data[static_cast<uint8_t>(addr_ptr[j])];
                        if (digit < 0) {
                                return false;
                        }

                        uint64_t hi;
                        const uint64_t tmp = num + umul128(order, static_cast<uint64_t>(digit), &hi);
                        if ((tmp < num) || hi) {
                                return false;
                        }

                        num = tmp;
                        order *= alphabet_size;
                }

                addr_ptr += (i < actual_num_full_blocks) ? block_sizes.back() : actual_last_block_size;  // Advance by actual size

                for (int j = static_cast<int>((i < actual_num_full_blocks) ? sizeof(num) : actual_last_block_size_index) - 1; j >= 0; --j) {
                        data[data_index++] = static_cast<uint8_t>(num >> (j * 8));
                }
        }

        // Decode varint tag from start of data
        uint64_t tag = 0;
        int varint_len = 0;
        for (int i = 0; i < 8; ++i) {
                tag |= static_cast<uint64_t>(data[i] & 0x7F) << (i * 7);
                ++varint_len;
                if ((data[i] & 0x80) == 0) {
                        break;
                }
        }

        char hex_buf[32];
        snprintf(hex_buf, sizeof(hex_buf), "0x%llx", static_cast<unsigned long long>(tag));

        m_prefix = tag;

        switch (m_prefix)
        {
        case valid_prefixes[0]: m_type = NetworkType::Mainnet;  break;
        case valid_prefixes[1]: m_type = NetworkType::Testnet;  break;
        case valid_prefixes[2]: m_type = NetworkType::Stagenet; break;

        case valid_prefixes_subaddress[0]: m_type = NetworkType::Mainnet;  m_subaddress = true; break;
        case valid_prefixes_subaddress[1]: m_type = NetworkType::Testnet;  m_subaddress = true; break;
        case valid_prefixes_subaddress[2]: m_type = NetworkType::Stagenet; m_subaddress = true; break;

        default:
                return false;
        }

        memcpy(m_spendPublicKey.h, data + varint_len, HASH_SIZE);
        memcpy(m_viewPublicKey.h, data + varint_len + HASH_SIZE, HASH_SIZE);

        // DEBUG: Print what we decoded
        {
            static constexpr char log_category_prefix[] = "Wallet ";
            char spend_hex[65] = {0}, view_hex[65] = {0}, data_hex[200] = {0}, prefix_hex[20] = {0};
            for (int i = 0; i < 32; i++) {
                sprintf(spend_hex + static_cast<ptrdiff_t>(i)*2, "%02x", m_spendPublicKey.h[i]);
                sprintf(view_hex + static_cast<ptrdiff_t>(i)*2, "%02x", m_viewPublicKey.h[i]);
            }
            for (int i = 0; i < std::min(data_index, 80); i++) {
                sprintf(data_hex + static_cast<ptrdiff_t>(i)*2, "%02x", data[i]);
            }
            sprintf(prefix_hex, "0x%llx", static_cast<unsigned long long>(m_prefix));
            LOGINFO(6, "decode varint_len=" << varint_len << " data_index=" << data_index << " prefix=" << static_cast<const char*>(prefix_hex));
            LOGINFO(6, "decode raw_data: " << static_cast<const char*>(data_hex));
            LOGINFO(6, "decode spend:    " << static_cast<const char*>(spend_hex));
            LOGINFO(6, "decode view:     " << static_cast<const char*>(view_hex));
        }

        // Load checksum from correct position (at end of decoded data)
        memcpy(&m_checksum, data + data_index - sizeof(m_checksum), sizeof(m_checksum));

        uint8_t md[200];
        keccak(data, static_cast<int>(data_index - sizeof(m_checksum)), md);

        uint32_t calculated_checksum;
        memcpy(&calculated_checksum, md, sizeof(calculated_checksum));

        if (m_checksum != calculated_checksum) {
                LOGINFO(1, "Checksum FAILED");
                m_type = NetworkType::Invalid;
        }

        if (memcmp(&m_checksum, md, sizeof(m_checksum)) != 0) {
                LOGINFO(1, "Checksum FAILED");
                m_type = NetworkType::Invalid;
        }

        ge_p3 point;
        if ((ge_frombytes_vartime(&point, m_spendPublicKey.h) != 0) || (ge_frombytes_vartime(&point, m_viewPublicKey.h) != 0)) {
                m_type = NetworkType::Invalid;
        }

        if (!torsion_check()) {
                LOGWARN(1, "Torsion check failed for wallet " << *this << "! It will not be compatible with FCMP++.");
        }

        return valid();
}

bool Wallet::assign(const hash& spend_pub_key, const hash& view_pub_key, NetworkType type, bool subaddress)
{
	ge_p3 point;
	if ((ge_frombytes_vartime(&point, spend_pub_key.h) != 0) || (ge_frombytes_vartime(&point, view_pub_key.h) != 0)) {
		return false;
	}

	switch (type)
	{
	case NetworkType::Mainnet:  m_prefix = subaddress ? valid_prefixes_subaddress[0] : valid_prefixes[0]; break;
	case NetworkType::Testnet:  m_prefix = subaddress ? valid_prefixes_subaddress[1] : valid_prefixes[1]; break;
	case NetworkType::Stagenet: m_prefix = subaddress ? valid_prefixes_subaddress[2] : valid_prefixes[2]; break;
	default:                    m_prefix = 0;                 break;
	}

	m_spendPublicKey = spend_pub_key;
	m_viewPublicKey = view_pub_key;

	uint8_t data[1 + HASH_SIZE * 2];
	data[0] = static_cast<uint8_t>(m_prefix);
	memcpy(data + 1, spend_pub_key.h, HASH_SIZE);
	memcpy(data + 1 + HASH_SIZE, view_pub_key.h, HASH_SIZE);

	uint8_t md[200];
	keccak(data, sizeof(data), md);

	memcpy(&m_checksum, md, sizeof(m_checksum));

	m_type = type;
	m_subaddress = subaddress;

	if (!torsion_check()) {
		LOGWARN(1, "Torsion check failed for wallet " << *this << "! It will not be compatible with FCMP++.");		
		// TODO: add "m_type = NetworkType::Invalid;" and return false in a later release, closer to FCMP++ hardfork
	}

	return true;
}

void Wallet::encode(char (&buf)[ADDRESS_LENGTH]) const
{
        uint8_t data[73]; // Max: 8 bytes varint + 32 spend + 32 view + 4 checksum
        int data_index = 0;
        
        // Write prefix as varint
        uint64_t prefix = m_prefix;
        do {
                data[data_index++] = static_cast<uint8_t>((prefix & 0x7F) | (prefix > 0x7F ? 0x80 : 0));
                prefix >>= 7;
        } while (prefix);
        
        // Write public keys
        memcpy(data + data_index, m_spendPublicKey.h, HASH_SIZE);
        memcpy(data + data_index + HASH_SIZE, m_viewPublicKey.h, HASH_SIZE);
        
        // Calculate and write checksum
        uint8_t md[200];
        const int pre_checksum_size = static_cast<int>(data_index + HASH_SIZE * 2);
        keccak(data, pre_checksum_size, md);
        memcpy(data + pre_checksum_size, md, sizeof(m_checksum));

        const int total_data_size = static_cast<int>(pre_checksum_size + sizeof(m_checksum));

        // Encode to base58 with variable-length data
        const int actual_num_full_blocks = static_cast<int>(total_data_size / sizeof(uint64_t));
        const int actual_last_block_bytes = static_cast<int>(total_data_size % sizeof(uint64_t));
        const int actual_last_block_size_index = actual_last_block_bytes > 0 ? block_sizes_lookup[actual_last_block_bytes] : -1;
        
        int buf_index = 0;
        for (int i = 0; i <= actual_num_full_blocks; ++i) {
                const bool is_last_block = (i == actual_num_full_blocks);
                const int bytes_in_block = is_last_block ? actual_last_block_bytes : static_cast<int>(sizeof(uint64_t));
                
                if (is_last_block && bytes_in_block == 0) break;
                
                // Read bytes in big-endian
                uint64_t n = 0;
                for (int j = 0; j < bytes_in_block; ++j) {
                        n = (n << 8) | data[i * sizeof(uint64_t) + j];
                }
                
                // Determine output block size
                const int output_block_size = is_last_block ? block_sizes[actual_last_block_size_index] : block_sizes.back();
                
                // Encode to base58
                for (int j = output_block_size - 1; j >= 0; --j) {
                        const int digit = static_cast<int>(n % alphabet_size);
                        n /= alphabet_size;
                        buf[buf_index + j] = alphabet[digit];
                }
                buf_index += output_block_size;
        }
        
        // Null terminate
        buf[buf_index] = '\0';
}

bool Wallet::get_eph_public_key(const hash& txkey_sec, size_t output_index, hash& eph_public_key, uint8_t& view_tag, const uint8_t* expected_view_tag) const
{
	hash derivation;
	if (!generate_key_derivation(m_viewPublicKey, txkey_sec, output_index, derivation, view_tag)) {
		return false;
	}

	if (expected_view_tag && (view_tag != *expected_view_tag)) {
		return false;
	}

	if (!derive_public_key(derivation, output_index, m_spendPublicKey, eph_public_key)) {
		return false;
	}

	return true;
}

bool Wallet::torsion_check() const
{
	ge_p3 p1, p2;
	if ((ge_frombytes_vartime(&p1, m_spendPublicKey.h) != 0) || (ge_frombytes_vartime(&p2, m_viewPublicKey.h) != 0)) {
		return false;
	}

	return
		!fcmp_pp::mul8_is_identity(p1) &&
		!fcmp_pp::mul8_is_identity(p2) &&
		fcmp_pp::torsion_check_vartime(p1) &&
		fcmp_pp::torsion_check_vartime(p2);
}

bool Wallet::get_eph_public_key_carrot(const hash& tx_key_seed, uint64_t height, size_t output_index, uint64_t amount, hash& eph_public_key, uint8_t& view_tag) const
{
    (void)output_index;  // Not used - anchor derived from spend pubkey, not position

    // 1. Build input_context from height
    uint8_t input_context[33];
    carrot::make_input_context_coinbase(height, input_context);

    // 2. Derive anchor from wallet's spend public key (position-independent)
    uint8_t anchor[16];
    carrot::derive_deterministic_anchor_from_pubkey(tx_key_seed, m_spendPublicKey, anchor);

    // 3. Derive per-output ephemeral private key d_e from anchor
    static const uint8_t null_payment_id[8] = {0};
    hash ephemeral_privkey;
    carrot::make_ephemeral_privkey(anchor, input_context, m_spendPublicKey, null_payment_id, ephemeral_privkey);

    // 4. Derive ephemeral public key D_e = d_e * B
    hash ephemeral_pubkey;
    carrot::make_ephemeral_pubkey_mainaddress(ephemeral_privkey, ephemeral_pubkey);

    // 5. Derive shared secret using this wallet's view pubkey
    hash shared_secret;
    if (!carrot::make_shared_secret_sender(ephemeral_privkey, m_viewPublicKey, shared_secret)) {
        return false;
    }

    // 6. Derive sender-receiver secret
    hash sender_receiver_secret;
    carrot::make_sender_receiver_secret(shared_secret, ephemeral_pubkey, input_context, sender_receiver_secret);

    // 7. Derive onetime address K_o
    carrot::make_onetime_address_coinbase(m_spendPublicKey, sender_receiver_secret, amount, eph_public_key);

    // 8. Derive view tag
    uint8_t view_tag_full[3];
    carrot::make_view_tag(shared_secret, input_context, eph_public_key, view_tag_full);
    view_tag = view_tag_full[0];

    return true;
}

} // namespace p2pool