#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/time.h>
#include <getopt.h> // コマンドライン引数処理用

// OpenCLヘッダー
#ifdef __APPLE__
#include <OpenCL/opencl.h>
#else
#include <CL/cl.h>
#endif

// --- 定義 ---

// SHA-256出力サイズ (256 bits = 32 bytes)
#define HASH_SIZE 32
// SHA-256ハッシュのワード数 (32ビットワード)
#define HASH_WORDS 8 
// SHA-256ブロックサイズ (バイト)
#define BLOCK_SIZE 64 

#define CHALLENGE_MAX_LEN 100
#define TARGET_HEX_MAX_LEN (HASH_SIZE * 2 + 1) // 64文字 + null終端
#define PREFIX "TKsaba_PoW_challange_"
#define PROGRESS_INTERVAL_MS 500 // 進捗表示の更新間隔 (ミリ秒)

// 進捗表示のための制御文字
#define CLEAR_LINE "\r\033[K"

// Nonceの型を cl_uint に固定
typedef cl_uint nonce_t;
// 32bitの全ビットセット = 4294967295。カーネルコードと合わせて uint の最大値を使用。
#define NONCE_NOT_FOUND ((nonce_t)-1) 


// --- ユーティリティ関数 ---

/**
 * @brief 16進数文字を値に変換
 * @param c 16進数文字
 * @return 対応する値 (0-15)
 */
static unsigned char hex_char_to_int(char c) {
    if (c >= '0' && c <= '9') return c - '0';
    if (c >= 'a' && c <= 'f') return c - 'a' + 10;
    if (c >= 'A' && c <= 'F') return c - 'A' + 10;
    return 0;
}

/**
 * @brief 16進数文字列をバイト配列に変換
 * @param hex_str 16進数文字列 (大文字/小文字)
 * @param byte_arr 格納先バイト配列 (サイズはstrlen(hex_str)/2)
 * @param max_len byte_arrの最大サイズ (バイト単位)
 * @return 変換されたバイト数。エラーの場合は-1。
 */
int hex_to_bytes(const char *hex_str, unsigned char *byte_arr, size_t max_len) {
    size_t len = strlen(hex_str);
    if (len % 2 != 0 || len / 2 > max_len) {
        return -1; // 長さが不正
    }

    size_t byte_len = len / 2;
    for (size_t i = 0; i < byte_len; i++) {
        unsigned char high = hex_char_to_int(hex_str[2 * i]);
        unsigned char low = hex_char_to_int(hex_str[2 * i + 1]);
        byte_arr[i] = (high << 4) | low;
    }
    return (int)byte_len;
}

/**
 * @brief バイト配列を16進数文字列に変換
 * @param byte_arr バイト配列
 * @param hex_str 格納先16進数文字列 (サイズはlen * 2 + 1)
 * @param len バイト配列の長さ
 */
void bytes_to_hex(const unsigned char *byte_arr, char *hex_str, size_t len) {
    for (size_t i = 0; i < len; i++) {
        sprintf(hex_str + (i * 2), "%02X", byte_arr[i]); // 大文字で表示
    }
    hex_str[len * 2] = '\0';
}

/**
 * @brief 高精度な現在時刻 (秒) を取得
 * @return 現在時刻 (秒)
 */
double get_current_time() {
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return tv.tv_sec + tv.tv_usec / 1000000.0;
}

// OpenCLエラーコードを文字列に変換する関数
const char *get_cl_error_string(cl_int err) {
    switch (err) {
        case CL_SUCCESS: return "CL_SUCCESS";
        case CL_DEVICE_NOT_FOUND: return "CL_DEVICE_NOT_FOUND";
        case CL_DEVICE_NOT_AVAILABLE: return "CL_DEVICE_NOT_AVAILABLE";
        case CL_COMPILER_NOT_AVAILABLE: return "CL_COMPILER_NOT_AVAILABLE";
        case CL_MEM_OBJECT_ALLOCATION_FAILURE: return "CL_MEM_OBJECT_ALLOCATION_FAILURE";
        case CL_OUT_OF_RESOURCES: return "CL_OUT_OF_RESOURCES";
        case CL_OUT_OF_HOST_MEMORY: return "CL_OUT_OF_HOST_MEMORY";
        case CL_PROFILING_INFO_NOT_AVAILABLE: return "CL_PROFILING_INFO_NOT_AVAILABLE";
        case CL_MEM_COPY_HOST_PTR: return "CL_MEM_COPY_HOST_PTR";
        case CL_INVALID_VALUE: return "CL_INVALID_VALUE";
        case CL_INVALID_DEVICE_TYPE: return "CL_INVALID_DEVICE_TYPE";
        case CL_INVALID_PLATFORM: return "CL_INVALID_PLATFORM";
        case CL_INVALID_DEVICE: return "CL_INVALID_DEVICE";
        case CL_INVALID_CONTEXT: return "CL_INVALID_CONTEXT";
        case CL_INVALID_QUEUE_PROPERTIES: return "CL_INVALID_QUEUE_PROPERTIES";
        case CL_INVALID_COMMAND_QUEUE: return "CL_INVALID_COMMAND_QUEUE";
        case CL_INVALID_HOST_PTR: return "CL_INVALID_HOST_PTR";
        case CL_INVALID_MEM_OBJECT: return "CL_INVALID_MEM_OBJECT";
        case CL_INVALID_KERNEL_ARGS: return "CL_INVALID_KERNEL_ARGS";
        case CL_INVALID_WORK_DIMENSION: return "CL_INVALID_WORK_DIMENSION";
        case CL_INVALID_WORK_GROUP_SIZE: return "CL_INVALID_WORK_GROUP_SIZE";
        case CL_INVALID_WORK_ITEM_SIZE: return "CL_INVALID_WORK_ITEM_SIZE";
        case CL_INVALID_GLOBAL_OFFSET: return "CL_INVALID_GLOBAL_OFFSET";
        case CL_INVALID_PROGRAM: return "CL_INVALID_PROGRAM";
        case CL_INVALID_PROGRAM_EXECUTABLE: return "CL_INVALID_PROGRAM_EXECUTABLE";
        case CL_INVALID_KERNEL: return "CL_INVALID_KERNEL";
        case CL_INVALID_EVENT_WAIT_LIST: return "CL_INVALID_EVENT_WAIT_LIST";
        case CL_INVALID_BUFFER_SIZE: return "CL_INVALID_BUFFER_SIZE";
        case CL_INVALID_GLOBAL_WORK_SIZE: return "CL_INVALID_GLOBAL_WORK_SIZE";
        case CL_BUILD_PROGRAM_FAILURE: return "CL_BUILD_PROGRAM_FAILURE";
        default: return "Unknown OpenCL Error";
    }
}

// --- OpenCLエラーチェックユーティリティ ---
void checkError(cl_int err, const char *name) {
    if (err != CL_SUCCESS) {
        fprintf(stderr, "ERROR: %s (%s)\n", name, get_cl_error_string(err));
        exit(1);
    }
}

// カーネルコードをファイルから読み込む関数
char* load_kernel_source(const char *filename) {
    FILE *file = fopen(filename, "r");
    if (!file) {
        perror("Failed to open kernel file: pow_solver.cl");
        return NULL;
    }
    
    // ファイルサイズの取得
    fseek(file, 0, SEEK_END);
    long size = ftell(file);
    rewind(file);

    // ソースコードの読み込み
    char *source = (char *)malloc(size + 1);
    if (!source) {
        perror("Failed to allocate memory for kernel source");
        fclose(file);
        return NULL;
    }
    
    if (fread(source, 1, size, file) != size) {
        perror("Failed to read kernel file");
        free(source);
        fclose(file);
        return NULL;
    }
    source[size] = '\0';
    fclose(file);
    return source;
}

/**
 * @brief ターゲットのバイト配列をOpenCLカーネルで扱いやすいuint[8]配列に変換 (ビッグエンディアン)
 */
void target_bytes_to_uints(const unsigned char *target_bytes, cl_uint *target_words) {
    // SHA-256はビッグエンディアン。バイト配列からビッグエンディアンのuintに変換
    for (int i = 0; i < HASH_WORDS; i++) {
        target_words[i] = (cl_uint)target_bytes[i * 4] << 24 |
                          (cl_uint)target_bytes[i * 4 + 1] << 16 |
                          (cl_uint)target_bytes[i * 4 + 2] << 8 |
                          (cl_uint)target_bytes[i * 4 + 3];
    }
}

// --- メイン関数 ---
int main(int argc, char *argv[]) {
    
    // --- 1. コマンドライン引数の解析 ---
    char *challenge_hex = NULL;
    char *target_hex_str = NULL;
    int opt;
    
    while ((opt = getopt(argc, argv, "c:t:")) != -1) {
        switch (opt) {
            case 'c':
                challenge_hex = optarg;
                break;
            case 't':
                target_hex_str = optarg;
                break;
            default:
                fprintf(stderr, "Usage: %s -c <challenge> -t <target_hash_hex>\n", argv[0]);
                return EXIT_FAILURE;
        }
    }

    // 必須引数のチェック
    if (!challenge_hex || !target_hex_str) {
        fprintf(stderr, "Error: Both challenge (-c) and target (-t) must be provided.\n");
        fprintf(stderr, "Usage: %s -c <challenge> -t <target_hash_hex>\n", argv[0]);
        return EXIT_FAILURE;
    }

    // --- 2. 初期化とパラメータ設定 ---
    unsigned char target_bytes[HASH_SIZE]; // 32バイト
    cl_uint target_words[HASH_WORDS]; // 8ワード (32バイト)
    char challenge_str[CHALLENGE_MAX_LEN];
    
    // チャレンジのコピーと検証
    if (strlen(challenge_hex) >= CHALLENGE_MAX_LEN) {
        fprintf(stderr, "Error: Challenge length exceeds %d characters.\n", CHALLENGE_MAX_LEN - 1);
        return EXIT_FAILURE;
    }
    strcpy(challenge_str, challenge_hex);

    // ターゲットの変換と検証
    if (hex_to_bytes(target_hex_str, target_bytes, HASH_SIZE) != HASH_SIZE) {
        fprintf(stderr, "Error: Target is not a valid 256-bit (64 character) hex string.\n");
        return EXIT_FAILURE;
    }
    
    // ターゲットを uint[8] に変換
    target_bytes_to_uints(target_bytes, target_words);
    
    // OpenCL変数の宣言
    cl_platform_id platform = NULL;
    cl_device_id device = NULL;
    cl_context context = NULL;
    cl_command_queue queue = NULL;
    cl_program program = NULL;
    cl_kernel kernel = NULL;
    cl_mem target_mem = NULL, base_block_mem = NULL, result_mem = NULL, final_hash_mem = NULL;
    cl_int err;
    
    // --- 3. OpenCL環境のセットアップ ---
    
    // a. プラットフォームとデバイスの取得
    err = clGetPlatformIDs(1, &platform, NULL);
    checkError(err, "clGetPlatformIDs");

    // デバイス（GPU）の取得を試みる
    err = clGetDeviceIDs(platform, CL_DEVICE_TYPE_GPU, 1, &device, NULL);
    if (err == CL_DEVICE_NOT_FOUND) {
        fprintf(stderr, "GPU not found, trying CPU...\n");
        // CPUの取得を試みる
        err = clGetDeviceIDs(platform, CL_DEVICE_TYPE_CPU, 1, &device, NULL);
    }
    checkError(err, "clGetDeviceIDs");

    // b. コンテキストとコマンドキューの作成
    context = clCreateContext(NULL, 1, &device, NULL, NULL, &err);
    checkError(err, "clCreateContext");
    // プロファイリングを有効にしてコマンドキューを作成
    queue = clCreateCommandQueue(context, device, CL_QUEUE_PROFILING_ENABLE, &err); 
    checkError(err, "clCreateCommandQueue");
    
    // c. カーネルのコンパイル
    char *kernel_source = load_kernel_source("pow_solver.cl");
    if (!kernel_source) return EXIT_FAILURE;
    
    program = clCreateProgramWithSource(context, 1, (const char **)&kernel_source, NULL, &err);
    checkError(err, "clCreateProgramWithSource");
    
    // OpenCL 1.2規格でビルド
    err = clBuildProgram(program, 1, &device, "-cl-std=CL1.2", NULL, NULL);
    if (err != CL_SUCCESS) {
        // ビルドログの取得と表示
        size_t log_size;
        clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, 0, NULL, &log_size);
        char *log = (char *)malloc(log_size);
        clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, log_size, log, NULL);
        fprintf(stderr, "OpenCL Build Error (clBuildProgram):\n%s\n", log);
        free(log);
        checkError(err, "clBuildProgram"); // エラーコードを出力して終了
    }
    free(kernel_source);
    
    kernel = clCreateKernel(program, "pow_hasher", &err);
    checkError(err, "clCreateKernel");
    
    // --- 4. データバッファの準備 ---
    
    // a. ベースブロックの作成
    char full_base_str[1024];
    strcpy(full_base_str, PREFIX);
    strcat(full_base_str, challenge_str);
    size_t base_len_chars = strlen(full_base_str); 
    // Nonce文字列の最大長（uintは最大10桁なので、NULL含め21桁あれば十分）
    cl_uint max_nonce_len = 21; 
    
    // 64バイトブロックのベースバッファを作成
    unsigned char base_block[64];
    memset(base_block, 0, 64);
    memcpy(base_block, full_base_str, base_len_chars);
    
    // Nonce文字列の開始位置 (バイト単位)
    cl_uint base_len = (cl_uint)base_len_chars; 
    
    // b. OpenCLバッファの作成
    
    // ターゲットハッシュ (読み取り専用)
    target_mem = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR,
                                sizeof(cl_uint) * HASH_WORDS, target_words, &err);
    checkError(err, "clCreateBuffer target_mem");

    // ベースブロック (Nonce埋め込み用)
    base_block_mem = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR,
                                    sizeof(base_block), base_block, &err);
    checkError(err, "clCreateBuffer base_block_mem");
    
    // 結果バッファ (Nonce、初期値は NONCE_NOT_FOUND)
    nonce_t found_nonce = NONCE_NOT_FOUND;
    result_mem = clCreateBuffer(context, CL_MEM_READ_WRITE | CL_MEM_COPY_HOST_PTR,
                                sizeof(nonce_t), &found_nonce, &err);
    checkError(err, "clCreateBuffer result_mem");
    
    // 最終ハッシュ
    cl_uint final_hash_words[HASH_WORDS];
    final_hash_mem = clCreateBuffer(context, CL_MEM_WRITE_ONLY,
                                    sizeof(cl_uint) * HASH_WORDS, NULL, &err);
    checkError(err, "clCreateBuffer final_hash_mem");

    // --- 5. カーネル引数の設定 (固定引数) ---
    // グローバルワークサイズ
    size_t global_work_size = 20000; // 2Mワークアイテム (調整可能)
    
    // Nonceのインクリメントはグローバルワークサイズに等しい
    nonce_t nonce_increment = (nonce_t)global_work_size; 
    
    err  = clSetKernelArg(kernel, 0, sizeof(cl_mem), &target_mem);
    err |= clSetKernelArg(kernel, 1, sizeof(cl_mem), &base_block_mem);
    // nonce_start はループ内で更新 (引数 2)
    err |= clSetKernelArg(kernel, 3, sizeof(nonce_t), &nonce_increment);
    err |= clSetKernelArg(kernel, 4, sizeof(cl_uint), &base_len);
    err |= clSetKernelArg(kernel, 5, sizeof(cl_uint), &max_nonce_len);
    err |= clSetKernelArg(kernel, 6, sizeof(cl_mem), &result_mem);
    err |= clSetKernelArg(kernel, 7, sizeof(cl_mem), &final_hash_mem);
    checkError(err, "clSetKernelArg fixed");

    // --- 6. 探索ループ ---
    nonce_t current_nonce_start = 0;
    long long total_hashes_tried = 0;
    double start_time = get_current_time();
    double last_progress_time = start_time;
    long long last_total_hashes = 0;
    
    // ローカルワークサイズ (最適値を設定するかNULLで自動設定)
    size_t local_work_size = 256; 
    
    // 探索の最大範囲 (uintの最大値)
    nonce_t max_nonce_value = NONCE_NOT_FOUND;

    fprintf(stdout, "--- OpenCL Proof-of-Work Solver ---\n");
    fprintf(stdout, "Challenge: %s\n", challenge_str);
    fprintf(stdout, "Target:    %s\n", target_hex_str);
    fprintf(stdout, "Searching with Global Work Size: %zu (Nonce range up to %u)\n", 
            global_work_size, max_nonce_value);

    while (found_nonce == NONCE_NOT_FOUND && current_nonce_start < max_nonce_value) {
        
        // Nonceの開始位置を設定 (引数 2)
        err = clSetKernelArg(kernel, 2, sizeof(nonce_t), &current_nonce_start);
        checkError(err, "clSetKernelArg nonce_start");
        
        // カーネル実行
        err = clEnqueueNDRangeKernel(queue, kernel, 1, NULL, &global_work_size, &local_work_size, 0, NULL, NULL);
        checkError(err, "clEnqueueNDRangeKernel");
        
        // コマンドキューをフラッシュし、GPUに作業を提出
        clFlush(queue);
        
        // 進捗報告と結果確認のために、少し待機
        usleep(PROGRESS_INTERVAL_MS * 1000); 
        
        // result_buffer から結果を読み戻し (完了を待機)
        err = clEnqueueReadBuffer(queue, result_mem, CL_TRUE, 0, sizeof(nonce_t), 
                                  &found_nonce, 0, NULL, NULL);
        checkError(err, "clEnqueueReadBuffer result_mem");
        
        // 進捗の更新
        total_hashes_tried += global_work_size;
        double current_time = get_current_time();
        
        if (current_time - last_progress_time >= (double)PROGRESS_INTERVAL_MS / 1000.0) {
             // 進捗表示
             double time_delta = current_time - last_progress_time;
             long long hashes_delta = total_hashes_tried - last_total_hashes;
             double hashrate = (time_delta > 0) ? (double)hashes_delta / time_delta : 0.0;
             fprintf(stderr, CLEAR_LINE "Progress: %lld hashes (%.2f H/s) | Searching from %u", 
                     total_hashes_tried, hashrate, current_nonce_start);
             fflush(stderr);
             last_total_hashes = total_hashes_tried;
             last_progress_time = current_time;
        }

        if (found_nonce != NONCE_NOT_FOUND) {
            // 見つかった
            break; 
        }

        // 次のNonce範囲へ
        current_nonce_start += nonce_increment;
        
        // オーバーフローチェック (uintの最大値に近づいた場合)
        if (current_nonce_start + nonce_increment < current_nonce_start) {
            // オーバーフローが発生した、または最大値を超えた
            printf("\nNonce search range reached maximum uint limit (%u).\n", max_nonce_value);
            break;
        }
    }
    
    // --- 7. 結果の取得と表示 ---
    clFinish(queue); // 全ての操作が完了するのを待機
    
    // 最終進捗行をクリア
    fprintf(stderr, CLEAR_LINE);
    
    double end_time = get_current_time();
    double time_elapsed = end_time - start_time;

    if (found_nonce != NONCE_NOT_FOUND) {
        // 最終ハッシュを読み戻し
        err = clEnqueueReadBuffer(queue, final_hash_mem, CL_TRUE, 0, 
                                  sizeof(cl_uint) * HASH_WORDS, final_hash_words, 0, NULL, NULL);
        checkError(err, "clEnqueueReadBuffer final_hash_mem");
        
        unsigned char final_hash_bytes[HASH_SIZE];
        char final_hash_hex[TARGET_HEX_MAX_LEN];
        
        // cl_uint[8]をバイト配列に変換
        for (int i = 0; i < HASH_WORDS; i++) {
            final_hash_bytes[i * 4]     = (unsigned char)(final_hash_words[i] >> 24);
            final_hash_bytes[i * 4 + 1] = (unsigned char)(final_hash_words[i] >> 16);
            final_hash_bytes[i * 4 + 2] = (unsigned char)(final_hash_words[i] >> 8);
            final_hash_bytes[i * 4 + 3] = (unsigned char)(final_hash_words[i]);
        }
        bytes_to_hex(final_hash_bytes, final_hash_hex, HASH_SIZE);

        fprintf(stdout, "\n--- SUCCESS (OpenCL) ---\n");
        fprintf(stdout, "Nonce found: %u\n", found_nonce);
        fprintf(stdout, "Final Hash:  %s\n", final_hash_hex); 
        fprintf(stdout, "Full Message: %s%u\n", full_base_str, found_nonce);
        fprintf(stdout, "Hashes tried (approx): %lld\n", total_hashes_tried);
        fprintf(stdout, "Time elapsed: %.4f sec\n", time_elapsed);
        fprintf(stdout, "Hash rate: %.2f H/s\n", (double)total_hashes_tried / time_elapsed);
    } else {
        printf("\nProof of Work not found within the 32-bit range (%u).\n", max_nonce_value);
    }

// --- 8. クリーンアップ ---
    if (target_mem) clReleaseMemObject(target_mem);
    if (base_block_mem) clReleaseMemObject(base_block_mem);
    if (result_mem) clReleaseMemObject(result_mem);
    if (final_hash_mem) clReleaseMemObject(final_hash_mem);
    if (kernel) clReleaseKernel(kernel);
    if (program) clReleaseProgram(program);
    if (queue) clReleaseCommandQueue(queue);
    if (context) clReleaseContext(context);

    return found_nonce != NONCE_NOT_FOUND ? EXIT_SUCCESS : EXIT_FAILURE;
}
