/**
 * main.cpp
 */

#include <iostream>
#include <getopt.h>
#include <string>
#include <netdb.h>
#include <cstring>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <vector>
#include <thread>
#include <iomanip>

#include "ee382n_bitcoin/stratum_client.hpp"
#include "ee382n_bitcoin/mining_work.hpp"
#include "ee382n_bitcoin/sha256.h"
#include "ee382n_bitcoin/bitcoin_block_header.hpp"

// Types
typedef union BITCOIN_BLOCK_HEADER_UNION {
    struct block_header_u8 {
        uint8_t version[4];
        uint8_t prev_block[32];
        uint8_t merkle_root[32];
        uint8_t timestamp[4];
        uint8_t difficulty_target[4];
        uint8_t nonce[4];
    } block_header_u8;
    struct block_header_u32 {
        uint32_t version[1];
        uint32_t prev_block[8];
        uint32_t merkle_root[8];
        uint32_t timestamp[1];
        uint32_t difficulty_target[1];
        uint32_t nonce[1];
    } block_header_u32;
    uint8_t bytes[80];
} bitcoin_block_header_t;

// Function declarations
int main(int argc, char *argv[]);
void on_new_work(bool clean_jobs);
void on_new_difficulty(double difficulty);
void mining_thread(const std::shared_ptr<stratum_v1::StratumClient>& client);
bool fulltest(const uint32_t *hash, const uint32_t *target);
void diff_to_target(uint32_t *target, double diff);

// Globals
bool get_new_work = false;
bool set_difficulty = false;


/**
 * Main
 * @param argc
 * @param argv
 * @return
 */
int main(int argc, char *argv[]) {
    // Parse command line arguments
    std::string pool, pool_url, worker_name, worker_pass;
    uint16_t pool_port;

    int c;
    while ((c = getopt(argc, argv, "o:u:p:")) != -1) {
        switch (c) {
            case 'o':
                pool = std::string(optarg);
                break;
            case 'u':
                worker_name = std::string(optarg);
                break;
            case 'p':
                worker_pass = std::string(optarg);
                break;
            case '?':
                if (optopt == 'o' || optopt == 'u' || optopt == 'p') {
                    std::cerr << "Option -" << optopt << " requires an argument." << std::endl;
                } else if (isprint(optopt)) {
                    std::cerr << "Unknown option -" << optopt << "." << std::endl;
                } else {
                    std::cerr << "Unknown option character `\\x`" << std::hex << optopt << "." << std::endl << std::dec;
                }
                return 1;
            default:
                abort();
        }
    }

    // Validate command line arguments
    if (pool.empty() || worker_name.empty() || worker_pass.empty()) {
        std::cerr << "Usage: " << argv[0] << " -o [pool] -u [worker_name] -p [worker_password]" << std::endl;
        std::cerr << "\t[pool]: Expected format: stratum+tcp://pool-ip:port" << std::endl;
        std::cerr << "\t[worker_name]: Name for worker" << std::endl;
        std::cerr << "\t[worker_password]: Any password" << std::endl;
        abort();
    } else {
        size_t found = pool.find_first_of(':');
        std::string protocol = pool.substr(0, found);
        std::string url_and_port = pool.substr(found + 3);
        size_t found1 = url_and_port.find_first_of(':');
        size_t found2 = url_and_port.find_first_of('/');
        pool_url = url_and_port.substr(0, found1);
        std::string port = found2 > found1 ? url_and_port.substr(found1 + 1, found2 - found1 - 1) : "";

        if (protocol != "stratum+tcp") {
            std::cerr << "Unknown protocol " << protocol << ". Only support stratum+tcp." << std::endl;
            abort();
        }

        if (pool_url.empty()) {
            std::cerr << "No pool URL provided." << std::endl;
            abort();
        }

        if (port.empty()) {
            std::cerr << "No port provided." << std::endl;
            abort();
        }
        pool_port = stoi(port);

        std::cout << "Mining pool connection parameters:" << std::endl;
        std::cout << "\tPROTOCOL = " << protocol << std::endl;
        std::cout << "\tURL = " << pool_url << std::endl;
        std::cout << "\tPORT = " << pool_port << std::endl;
        std::cout << "\tWORKER_NAME = " << worker_name << std::endl;
        std::cout << "\tWORKER_PASSWORD = " << worker_pass << std::endl;
    }

    std::cout << "TEST" << std::endl;

    BitcoinBlockHeader h1 = BitcoinBlockHeader::GenesisBlock();
    assert(h1.hash() == "000000000019d6689c085ae165831e934ff763ae46a2a6c172b3f1b60a8ce26f");

    std::cout << h1.hash() << std::endl;

    BitcoinBlockHeader h2;
    h2.set_version(536870912);
    h2.set_prev_block("000000000029b6437babf906ae1ac0c1b30b6ae97a83bb80c79f9b5cf1650a73");
    h2.set_merkle_root("c8171521630dbb68165b3524538ebb17dcc0f942b33b5bd19fe58b1dd6c26931");
    h2.set_ntime(1713725552);
    h2.set_nbits(422451157);
    h2.set_nonce(83899);
    assert(h2.hash() == "8e347681059365e8d00f0903f1399d2bb92c3b4e9f28edb0aaa4def03c36e223");

    std::cout << h2.hash() << std::endl;

    std::cout << "END TEST" << std::endl;

    std::cout << "Connecting to mining pool " << pool << std::endl;

    auto client = std::make_shared<stratum_v1::StratumClient>(worker_name, worker_pass);
    client->connect(pool_url, pool_port);

    // Authorize with pool to begin receiving work
    client->mining_subscribe();
    client->mining_authorize();

    client->new_work_callback = std::function<void(bool)>(on_new_work);
    client->set_difficulty_callback = std::function<void(double)>(on_new_difficulty);
    std::thread cpu_miner(mining_thread, client);

    client->spin_forever();

    return 0;
}

/**
 * Callback triggered when server sends new job to the work queue
 * @param clean_jobs If true, clear queue and drop any old work, start on new job
 */
void on_new_work(bool clean_jobs) {
    if(clean_jobs) {
        std::cout << "[on_new_work] Sending signal to restart miner" << std::endl;
        get_new_work = true;
    }
}

void on_new_difficulty(double difficulty) {
    set_difficulty = true;
}

/**
 * Single-threaded CPU miner
 */
void mining_thread(const std::shared_ptr<stratum_v1::StratumClient>& client) {
    // Statistics
    unsigned long hash_count = 0;
    unsigned long valid_count = 0;
    unsigned long accepted_count = 0; // TODO: Track accepted share metrics
    clock_t begin = clock();

    while(1) {
        // Get job from job queue
        stratum_v1::MiningWork work = client->get_work();
        BitcoinBlockHeader block_header = work.generate_header_template();
        uint32_t nonce = 0;
        uint8_t block_hash[32];
        uint32_t target[8];
        diff_to_target(target, client->get_difficulty());
        get_new_work = false; // Reset flag

        std::cout << "[mining_thread] Starting new job" << std::endl;

        // TODO: Iterate over 16 version bits (header) and xnonce2 (coin base)
        do {
            block_header.set_nonce(nonce);
            sha256_double(block_header.header_bytes(), 80, block_hash);

            if(set_difficulty) {
                diff_to_target(target, client->get_difficulty());
                set_difficulty = false;
            }

            if((hash_count++ % 100000) == 0) {
                clock_t now = clock();
                double rate = (double)hash_count / ((double)(now - begin) / CLOCKS_PER_SEC);
                std::cout << "[mining_thread] hash_count: " << hash_count << " (" << rate << " H/s) valid_count: "
                          << valid_count << " accepted_count: " << accepted_count << " invalid_count: "
                          << (valid_count - accepted_count) << std::endl;
            }

            // TODO: Implement fast test, only run fulltest if fast test passes?
            if(fulltest(reinterpret_cast<uint32_t*>(block_hash), target)) {
                client->mining_submit(work, nonce);
                valid_count++;
                std::cout << "[mining_thread] Found valid share!" << std::endl;
                std::cout << "\tHeader:\t" << block_header << std::endl;
                std::cout << "\tHash:\t" << stratum_v1::MiningWork::bin2hex(block_hash, 32) << std::endl;

                break;
            }

            nonce++;
        } while(nonce != 0 && !get_new_work);

        if(get_new_work) {
            std::cout << "[mining_thread] Stopped early to get new work" << std::endl;
        }
    }
}

/**
 * fulltest: Check if 256 bit hash is valid
 * https://github.com/pooler/cpuminer/blob/5f02105940edb61144c09a7eb960bba04a10d5b7/util.c#L839
 *
 * @param hash 256 bit hash
 * @param target 256 bit difficulty target
 * @return True if hash is valid
 */
bool fulltest(const uint32_t *hash, const uint32_t *target) {
    int i;
    bool rc = true;

    for (i = 7; i >= 0; i--) {
        if (hash[i] > target[i]) {
            rc = false;
            break;
        }
        if (hash[i] < target[i]) {
            rc = true;
            break;
        }
    }

    return rc;
}

/**
 * diff_to_target converts pool floating point difficulty to 256 bit difficulty target
 * @param target Target from pool
 * @param diff Difficulty
 */
void diff_to_target(uint32_t *target, double diff)
{
    uint64_t m;
    int k;

    for (k = 6; k > 0 && diff > 1.0; k--)
        diff /= 4294967296.0;
    m = 4294901760.0 / diff;
    if (m == 0 && k == 6)
        memset(target, 0xff, 32);
    else {
        memset(target, 0, 32);
        target[k] = (uint32_t)m;
        target[k + 1] = (uint32_t)(m >> 32);
    }
}