#include <poll.h>
#include <sched.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

#include "config.h"
#include "mining_loop.h"
#include "rpc.h"
#include "utils.h"

typedef struct {
    pid_t pid;
    int fd;
    char buf[4096];
    size_t len;
} WorkerProc;

typedef struct {
    double *rates;
    long long *attempts;
    char block_hash[130];
    int winner_idx;
    double winner_rate;
    int has_winner_rate;
    int lines_printed;
    double start_t;
    double last_print;
} DashboardState;

static double now_seconds(void) {
    struct timespec ts;
    clock_gettime(CLOCK_REALTIME, &ts);
    return (double)ts.tv_sec + (double)ts.tv_nsec / 1e9;
}

static void clear_lines(int n) {
    int i;
    for (i = 0; i < n; i++) {
        fputs("\033[F\033[K", stdout);
    }
    fflush(stdout);
}

static void dashboard_print(DashboardState *st, int n) {
    int i;
    double total_rate = 0.0;
    long long total_attempts = 0;
    time_t now_t = time(NULL);
    char ts[64];

    if (st->lines_printed > 0) {
        clear_lines(st->lines_printed);
    }

    for (i = 0; i < n; i++) {
        total_rate += st->rates[i];
        total_attempts += st->attempts[i];
    }

    strftime(ts, sizeof(ts), "%Y-%m-%d %H:%M:%S", localtime(&now_t));

    printf("%s | MINING STATUS\n", ts);
    printf("========================================\n");
    printf("Total: %.2f kH/s | Attempts: %lld\n", total_rate, total_attempts);
    printf("----------------------------------------\n");
    for (i = 0; i < n; i++) {
        printf("Worker %-2d: %.2f kH/s | Attempts: %lld\n", i, st->rates[i], st->attempts[i]);
    }

    st->lines_printed = 4 + n;
    fflush(stdout);
}

static int parse_event_line(const char *line, DashboardState *st, int n) {
    int idx;
    double rate;
    long long attempts;
    char hash[129];

    if (sscanf(line, "status %d %lf %lld", &idx, &rate, &attempts) == 3) {
        if (idx >= 0 && idx < n) {
            st->rates[idx] = rate;
            st->attempts[idx] = attempts;
        }
        return 0;
    }

    if (sscanf(line, "found %d %lf", &idx, &rate) == 2) {
        st->winner_idx = idx;
        st->winner_rate = rate;
        st->has_winner_rate = 1;
        return 0;
    }

    if (sscanf(line, "hash %d %128s", &idx, hash) == 2) {
        snprintf(st->block_hash, sizeof(st->block_hash), "%s", hash);
        return 0;
    }

    if (sscanf(line, "submit %d", &idx) == 1) {
        return 1;
    }

    return 0;
}

static int process_worker_bytes(WorkerProc *w, DashboardState *st, int n) {
    char tmp[512];
    ssize_t r;

    r = read(w->fd, tmp, sizeof(tmp));
    if (r < 0) {
        return 0;
    }

    if (r == 0) {
        close(w->fd);
        w->fd = -1;
        return 0;
    }

    if (w->len + (size_t)r >= sizeof(w->buf)) {
        w->len = 0;
    }

    memcpy(w->buf + w->len, tmp, (size_t)r);
    w->len += (size_t)r;

    while (1) {
        char *nl = memchr(w->buf, '\n', w->len);
        size_t line_len;
        char line[512];
        int rc;

        if (nl == NULL) {
            break;
        }

        line_len = (size_t)(nl - w->buf);
        if (line_len >= sizeof(line)) {
            line_len = sizeof(line) - 1;
        }

        memcpy(line, w->buf, line_len);
        line[line_len] = '\0';

        memmove(w->buf, nl + 1, w->len - (line_len + 1));
        w->len -= (line_len + 1);

        rc = parse_event_line(line, st, n);
        if (rc == 1) {
            return 1;
        }
    }

    return 0;
}

static int aggregate_loop(WorkerProc *workers, int n) {
    struct pollfd *pfds;
    DashboardState st;
    int i;

    memset(&st, 0, sizeof(st));
    st.rates = (double *)calloc((size_t)n, sizeof(double));
    st.attempts = (long long *)calloc((size_t)n, sizeof(long long));
    st.winner_idx = -1;
    st.start_t = now_seconds();
    st.last_print = 0.0;

    if (st.rates == NULL || st.attempts == NULL) {
        free(st.rates);
        free(st.attempts);
        return 0;
    }

    pfds = (struct pollfd *)calloc((size_t)n, sizeof(struct pollfd));
    if (pfds == NULL) {
        free(st.rates);
        free(st.attempts);
        return 0;
    }

    while (1) {
        int active = 0;
        double now = now_seconds();
        int submit_seen = 0;

        for (i = 0; i < n; i++) {
            pfds[i].fd = workers[i].fd;
            pfds[i].events = (workers[i].fd >= 0) ? POLLIN : 0;
            pfds[i].revents = 0;
            if (workers[i].fd >= 0) {
                active++;
            }
        }

        if (active == 0) {
            break;
        }

        poll(pfds, (nfds_t)n, 200);

        for (i = 0; i < n; i++) {
            if (workers[i].fd >= 0 && (pfds[i].revents & POLLIN)) {
                if (process_worker_bytes(&workers[i], &st, n)) {
                    submit_seen = 1;
                }
            }
            if (workers[i].fd >= 0 && (pfds[i].revents & (POLLHUP | POLLERR | POLLNVAL))) {
                close(workers[i].fd);
                workers[i].fd = -1;
            }
        }

        if (now - st.last_print >= 1.0) {
            dashboard_print(&st, n);
            st.last_print = now;
        }

        if (submit_seen) {
            double elapsed = now_seconds() - st.start_t;
            long long total_attempts = 0;
            double avg_rate;

            if (st.lines_printed > 0) {
                clear_lines(st.lines_printed);
            }

            for (i = 0; i < n; i++) {
                total_attempts += st.attempts[i];
            }

            avg_rate = (elapsed > 0.0) ? (double)total_attempts / elapsed / 1000.0 : 0.0;

            printf("==============================================================================\n");
            printf("[OK] BLOCK FOUND AND SUBMITTED\n");
            printf("  Hash: %s\n", st.block_hash[0] ? st.block_hash : "N/A");
            if (st.winner_idx >= 0) {
                printf("  Worker: %d\n", st.winner_idx);
            }
            if (st.has_winner_rate) {
                printf("  Worker hashrate: %.2f kH/s\n", st.winner_rate);
            }
            printf("  Average total hashrate: %.2f kH/s\n", avg_rate);
            printf("  Total attempts: %lld\n", total_attempts);
            printf("==============================================================================\n");

            free(pfds);
            free(st.rates);
            free(st.attempts);
            return 1;
        }
    }

    free(pfds);
    free(st.rates);
    free(st.attempts);
    return 0;
}

static void terminate_workers(WorkerProc *workers, int n) {
    int i;

    for (i = 0; i < n; i++) {
        if (workers[i].pid > 0) {
            kill(workers[i].pid, SIGTERM);
        }
    }

    for (i = 0; i < n; i++) {
        if (workers[i].pid > 0) {
            waitpid(workers[i].pid, NULL, 0);
        }
        if (workers[i].fd >= 0) {
            close(workers[i].fd);
            workers[i].fd = -1;
        }
    }
}

static int spawn_worker(WorkerProc *w, int idx, const MinerConfig *cfg, const char *base_ex2) {
    int p[2];
    pid_t pid;

    if (pipe(p) < 0) {
        return 0;
    }

    pid = fork();
    if (pid < 0) {
        close(p[0]);
        close(p[1]);
        return 0;
    }

    if (pid == 0) {
        char *ex2;
        close(p[0]);

#ifdef __linux__
        {
            cpu_set_t set;
            CPU_ZERO(&set);
            CPU_SET(idx % (int)sysconf(_SC_NPROCESSORS_ONLN), &set);
            sched_setaffinity(0, sizeof(set), &set);
        }
#endif

        ex2 = hex_add_width(base_ex2, idx);
        if (ex2 == NULL) {
            _exit(1);
        }

        srand((unsigned int)(time(NULL) ^ getpid()));
        run_mining_loop(cfg, idx, ex2, p[1]);
        free(ex2);
        close(p[1]);
        _exit(0);
    }

    close(p[1]);

    w->pid = pid;
    w->fd = p[0];
    w->len = 0;
    return 1;
}

static void usage(const char *prog) {
    fprintf(stderr, "Uso: %s [--config miner.conf] [-n NUM_PROCS] [--base-extranonce2 HEX]\n", prog);
}

int main(int argc, char **argv) {
    const char *config_path = "miner.conf";
    int num_procs = -1;
    const char *base_ex2_cli = NULL;
    MinerConfig cfg;
    RpcClient rpc;
    char *chain = NULL;
    char *base_ex2;
    int i;

    for (i = 1; i < argc; i++) {
        if ((strcmp(argv[i], "-n") == 0 || strcmp(argv[i], "--num-procs") == 0) && i + 1 < argc) {
            num_procs = atoi(argv[++i]);
            continue;
        }
        if (strcmp(argv[i], "--base-extranonce2") == 0 && i + 1 < argc) {
            base_ex2_cli = argv[++i];
            continue;
        }
        if (strcmp(argv[i], "--config") == 0 && i + 1 < argc) {
            config_path = argv[++i];
            continue;
        }
        usage(argv[0]);
        return 1;
    }

    if (config_load(config_path, &cfg) == 0) {
        return 1;
    }

    if (num_procs <= 0) {
        num_procs = cfg.num_processors;
    }

    base_ex2 = (char *)(base_ex2_cli ? base_ex2_cli : cfg.extranonce2);

    rpc_init(&rpc, &cfg);
    if (rpc_test_connection(&rpc, &chain) == 0) {
        fprintf(stderr, "[launcher] RPC test fallito\n");
        return 1;
    }

    printf("\nStarting mining with %d processes (chain=%s, extranonce2 base=%s)\n\n", num_procs, chain, base_ex2);
    free(chain);

    while (1) {
        WorkerProc *workers = (WorkerProc *)calloc((size_t)num_procs, sizeof(WorkerProc));
        int ok = 1;
        int restart;

        if (workers == NULL) {
            return 1;
        }

        for (i = 0; i < num_procs; i++) {
            if (spawn_worker(&workers[i], i, &cfg, base_ex2) == 0) {
                ok = 0;
                break;
            }
        }

        if (ok == 0) {
            fprintf(stderr, "[launcher] errore nello spawn worker\n");
            terminate_workers(workers, num_procs);
            free(workers);
            return 1;
        }

        restart = aggregate_loop(workers, num_procs);
        terminate_workers(workers, num_procs);
        free(workers);

        if (restart == 0) {
            break;
        }

        printf("\nRestarting workers...\n\n");
        sleep(1);
    }

    return 0;
}