#include "block_builder.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "utils.h"

static int append_u32_le_hex(StrBuf *sb, uint32_t v) {
    return sb_append_byte_hex(sb, (uint8_t)(v & 0xFF)) &&
           sb_append_byte_hex(sb, (uint8_t)((v >> 8) & 0xFF)) &&
           sb_append_byte_hex(sb, (uint8_t)((v >> 16) & 0xFF)) &&
           sb_append_byte_hex(sb, (uint8_t)((v >> 24) & 0xFF));
}

static int append_u64_le_hex(StrBuf *sb, uint64_t v) {
    int i;
    for (i = 0; i < 8; i++) {
        if (sb_append_byte_hex(sb, (uint8_t)((v >> (8 * i)) & 0xFF)) == 0) {
            return 0;
        }
    }
    return 1;
}

static char *ascii_to_hex(const char *s) {
    StrBuf sb;
    size_t i;

    sb_init(&sb);
    for (i = 0; s[i] != '\0'; ++i) {
        if (sb_append_byte_hex(&sb, (uint8_t)s[i]) == 0) {
            sb_free(&sb);
            return NULL;
        }
    }

    return sb_take(&sb);
}

char *tx_encode_coinbase_height(int height) {
    uint8_t tmp[8];
    int n = 0;
    int v;
    StrBuf sb;
    int i;

    if (height < 0) {
        return NULL;
    }

    if (height == 0) {
        return strdup("00");
    }

    v = height;
    while (v > 0) {
        tmp[n++] = (uint8_t)(v & 0xFF);
        v >>= 8;
    }

    if (tmp[n - 1] & 0x80) {
        tmp[n++] = 0x00;
    }

    sb_init(&sb);
    if (sb_append_byte_hex(&sb, (uint8_t)n) == 0) {
        sb_free(&sb);
        return NULL;
    }

    for (i = 0; i < n; i++) {
        if (sb_append_byte_hex(&sb, tmp[i]) == 0) {
            sb_free(&sb);
            return NULL;
        }
    }

    return sb_take(&sb);
}

int is_segwit_tx(const char *raw_hex) {
    size_t n;

    if (raw_hex == NULL) {
        return 0;
    }

    n = strlen(raw_hex);
    if (n < 12) {
        return 0;
    }

    return raw_hex[8] == '0' && raw_hex[9] == '0' && raw_hex[10] == '0' && raw_hex[11] == '1';
}

int build_coinbase_transaction(
    const BlockTemplate *tpl,
    const char *miner_script_pubkey,
    const char *extranonce1,
    const char *extranonce2,
    const char *coinbase_message,
    char **coinbase_hex_out,
    char **coinbase_txid_out
) {
    int segwit;
    StrBuf script_sig;
    StrBuf tx;
    StrBuf outputs;
    char *height_hex;
    char *msg_hex = NULL;
    char *coinbase_hex;
    char *core = NULL;
    uint8_t *core_bytes = NULL;
    size_t core_len = 0;
    uint8_t txid_bin[32];
    uint8_t txid_rev[32];

    *coinbase_hex_out = NULL;
    *coinbase_txid_out = NULL;

    segwit = (tpl->default_witness_commitment != NULL && tpl->default_witness_commitment[0] != '\0');

    height_hex = tx_encode_coinbase_height(tpl->height);
    if (height_hex == NULL) {
        return 0;
    }

    sb_init(&script_sig);
    if (sb_append(&script_sig, height_hex) == 0) {
        free(height_hex);
        sb_free(&script_sig);
        return 0;
    }
    free(height_hex);

    if (coinbase_message != NULL && coinbase_message[0] != '\0') {
        size_t mlen = strlen(coinbase_message);
        msg_hex = ascii_to_hex(coinbase_message);
        if (msg_hex == NULL) {
            sb_free(&script_sig);
            return 0;
        }
        if (sb_append(&script_sig, "6a") == 0 ||
            sb_append_byte_hex(&script_sig, (uint8_t)mlen) == 0 ||
            sb_append(&script_sig, msg_hex) == 0) {
            free(msg_hex);
            sb_free(&script_sig);
            return 0;
        }
        free(msg_hex);
    }

    if (sb_append(&script_sig, extranonce1) == 0 || sb_append(&script_sig, extranonce2) == 0) {
        sb_free(&script_sig);
        return 0;
    }

    if ((script_sig.len / 2) > 100) {
        fprintf(stderr, "[builder] scriptSig > 100 bytes\n");
        sb_free(&script_sig);
        return 0;
    }

    sb_init(&tx);
    if (append_u32_le_hex(&tx, 2) == 0) {
        sb_free(&script_sig);
        sb_free(&tx);
        return 0;
    }

    if (segwit) {
        if (sb_append(&tx, "0001") == 0) {
            sb_free(&script_sig);
            sb_free(&tx);
            return 0;
        }
    }

    if (sb_append(&tx, "01") == 0 ||
        sb_append(&tx, "0000000000000000000000000000000000000000000000000000000000000000") == 0 ||
        sb_append(&tx, "ffffffff") == 0) {
        sb_free(&script_sig);
        sb_free(&tx);
        return 0;
    }

    if (encode_varint_hex((uint64_t)(script_sig.len / 2), &tx) == 0 || sb_append(&tx, script_sig.data) == 0) {
        sb_free(&script_sig);
        sb_free(&tx);
        return 0;
    }

    if (sb_append(&tx, "ffffffff") == 0) {
        sb_free(&script_sig);
        sb_free(&tx);
        return 0;
    }

    sb_free(&script_sig);

    sb_init(&outputs);

    if (append_u64_le_hex(&outputs, tpl->coinbase_value) == 0 ||
        encode_varint_hex((uint64_t)(strlen(miner_script_pubkey) / 2), &outputs) == 0 ||
        sb_append(&outputs, miner_script_pubkey) == 0) {
        sb_free(&outputs);
        sb_free(&tx);
        return 0;
    }

    if (segwit) {
        const char *wc_raw = tpl->default_witness_commitment;
        StrBuf wc_script;

        sb_init(&wc_script);
        if (wc_raw[0] == '6' && wc_raw[1] == 'a') {
            if (sb_append(&wc_script, wc_raw) == 0) {
                sb_free(&wc_script);
                sb_free(&outputs);
                sb_free(&tx);
                return 0;
            }
        } else {
            if (sb_append(&wc_script, "6a24aa21a9ed") == 0 || sb_append(&wc_script, wc_raw) == 0) {
                sb_free(&wc_script);
                sb_free(&outputs);
                sb_free(&tx);
                return 0;
            }
        }

        if (append_u64_le_hex(&outputs, 0) == 0 ||
            encode_varint_hex((uint64_t)(wc_script.len / 2), &outputs) == 0 ||
            sb_append(&outputs, wc_script.data) == 0) {
            sb_free(&wc_script);
            sb_free(&outputs);
            sb_free(&tx);
            return 0;
        }

        sb_free(&wc_script);
    }

    if (encode_varint_hex(segwit ? 2 : 1, &tx) == 0 || sb_append(&tx, outputs.data) == 0) {
        sb_free(&outputs);
        sb_free(&tx);
        return 0;
    }

    sb_free(&outputs);

    if (segwit) {
        if (sb_append(&tx, "0120") == 0 ||
            sb_append(&tx, "0000000000000000000000000000000000000000000000000000000000000000") == 0) {
            sb_free(&tx);
            return 0;
        }
    }

    if (sb_append(&tx, "00000000") == 0) {
        sb_free(&tx);
        return 0;
    }

    coinbase_hex = sb_take(&tx);

    if (segwit) {
        size_t core_len_hex;
        size_t lock_start;
        size_t body_len;
        size_t no_wit_len;

        core_len_hex = 8 + (strlen(coinbase_hex) - 12);
        core = (char *)malloc(core_len_hex + 1);
        if (core == NULL) {
            free(coinbase_hex);
            return 0;
        }

        memcpy(core, coinbase_hex, 8);
        memcpy(core + 8, coinbase_hex + 12, strlen(coinbase_hex) - 12 + 1);

        lock_start = strlen(core) - 8;
        body_len = lock_start;
        if (body_len < 68) {
            free(core);
            free(coinbase_hex);
            return 0;
        }
        no_wit_len = body_len - 68;

        memmove(core + no_wit_len, core + lock_start, 8);
        core[no_wit_len + 8] = '\0';
    } else {
        core = strdup(coinbase_hex);
        if (core == NULL) {
            free(coinbase_hex);
            return 0;
        }
    }

    if (hex_to_bytes(core, &core_bytes, &core_len) == 0) {
        free(core);
        free(coinbase_hex);
        return 0;
    }

    double_sha256(core_bytes, core_len, txid_bin);
    memcpy(txid_rev, txid_bin, 32);
    reverse_bytes(txid_rev, 32);

    *coinbase_txid_out = bytes_to_hex(txid_rev, 32);
    *coinbase_hex_out = coinbase_hex;

    free(core_bytes);
    free(core);

    if (*coinbase_txid_out == NULL) {
        free(*coinbase_hex_out);
        *coinbase_hex_out = NULL;
        return 0;
    }

    return 1;
}

char *calculate_merkle_root(const char *coinbase_txid_hex, const TemplateTransaction *txs, size_t tx_count) {
    size_t leaf_count = tx_count + 1;
    uint8_t *nodes;
    size_t level_count;
    size_t i;

    nodes = (uint8_t *)malloc(leaf_count * 32);
    if (nodes == NULL) {
        return NULL;
    }

    {
        uint8_t coinbase[32];
        if (hex_to_fixed_bytes(coinbase_txid_hex, coinbase, 32) == 0) {
            free(nodes);
            return NULL;
        }
        reverse_bytes(coinbase, 32);
        memcpy(nodes, coinbase, 32);
    }

    for (i = 0; i < tx_count; i++) {
        uint8_t h[32];
        if (hex_to_fixed_bytes(txs[i].hash, h, 32) == 0) {
            free(nodes);
            return NULL;
        }
        reverse_bytes(h, 32);
        memcpy(nodes + (i + 1) * 32, h, 32);
    }

    level_count = leaf_count;
    while (level_count > 1) {
        size_t next_count = (level_count + 1) / 2;
        uint8_t *next_nodes = (uint8_t *)malloc(next_count * 32);

        if (next_nodes == NULL) {
            free(nodes);
            return NULL;
        }

        for (i = 0; i < next_count; i++) {
            uint8_t pair[64];
            uint8_t dh[32];
            size_t left = i * 2;
            size_t right = left + 1;

            memcpy(pair, nodes + left * 32, 32);
            if (right < level_count) {
                memcpy(pair + 32, nodes + right * 32, 32);
            } else {
                memcpy(pair + 32, nodes + left * 32, 32);
            }

            double_sha256(pair, 64, dh);
            memcpy(next_nodes + i * 32, dh, 32);
        }

        free(nodes);
        nodes = next_nodes;
        level_count = next_count;
    }

    reverse_bytes(nodes, 32);
    {
        char *root = bytes_to_hex(nodes, 32);
        free(nodes);
        return root;
    }
}

int build_block_header_bytes(
    int version,
    const char *prev_hash_hex,
    const char *merkle_root_hex,
    uint32_t timestamp,
    const char *bits_hex,
    uint32_t nonce,
    uint8_t out_header[80]
) {
    uint8_t prev[32];
    uint8_t merkle[32];
    uint8_t bits[4];

    if (hex_to_fixed_bytes(prev_hash_hex, prev, 32) == 0) {
        return 0;
    }
    if (hex_to_fixed_bytes(merkle_root_hex, merkle, 32) == 0) {
        return 0;
    }
    if (hex_to_fixed_bytes(bits_hex, bits, 4) == 0) {
        return 0;
    }

    reverse_bytes(prev, 32);
    reverse_bytes(merkle, 32);
    reverse_bytes(bits, 4);

    out_header[0] = (uint8_t)(version & 0xFF);
    out_header[1] = (uint8_t)((version >> 8) & 0xFF);
    out_header[2] = (uint8_t)((version >> 16) & 0xFF);
    out_header[3] = (uint8_t)((version >> 24) & 0xFF);

    memcpy(out_header + 4, prev, 32);
    memcpy(out_header + 36, merkle, 32);

    out_header[68] = (uint8_t)(timestamp & 0xFF);
    out_header[69] = (uint8_t)((timestamp >> 8) & 0xFF);
    out_header[70] = (uint8_t)((timestamp >> 16) & 0xFF);
    out_header[71] = (uint8_t)((timestamp >> 24) & 0xFF);

    memcpy(out_header + 72, bits, 4);

    out_header[76] = (uint8_t)(nonce & 0xFF);
    out_header[77] = (uint8_t)((nonce >> 8) & 0xFF);
    out_header[78] = (uint8_t)((nonce >> 16) & 0xFF);
    out_header[79] = (uint8_t)((nonce >> 24) & 0xFF);

    return 1;
}

char *build_block_header_hex(
    int version,
    const char *prev_hash_hex,
    const char *merkle_root_hex,
    uint32_t timestamp,
    const char *bits_hex,
    uint32_t nonce
) {
    uint8_t header[80];

    if (build_block_header_bytes(version, prev_hash_hex, merkle_root_hex, timestamp, bits_hex, nonce, header) == 0) {
        return NULL;
    }

    return bytes_to_hex(header, 80);
}

char *serialize_block(const char *header_hex, const char *coinbase_tx_hex, const TemplateTransaction *txs, size_t tx_count) {
    StrBuf sb;
    size_t i;

    sb_init(&sb);

    if (sb_append(&sb, header_hex) == 0) {
        sb_free(&sb);
        return NULL;
    }

    if (encode_varint_hex((uint64_t)(tx_count + 1), &sb) == 0) {
        sb_free(&sb);
        return NULL;
    }

    if (sb_append(&sb, coinbase_tx_hex) == 0) {
        sb_free(&sb);
        return NULL;
    }

    for (i = 0; i < tx_count; i++) {
        if (txs[i].data == NULL || sb_append(&sb, txs[i].data) == 0) {
            sb_free(&sb);
            return NULL;
        }
    }

    return sb_take(&sb);
}