#include "common.h"
#include <sys/sendfile.h>
#include <getopt.h>
#include <dirent.h>
#include <stdatomic.h>
#include <time.h>
#include <netinet/tcp.h>
#include <stdio.h>

// ── Helper to copy a file ────────────────────────────────────────────────
static int copy_file(const char *src, const char *dst) {
    int src_fd = open(src, O_RDONLY);
    if (src_fd < 0) return -1;
    
    int dst_fd = open(dst, O_WRONLY | O_CREAT | O_TRUNC, 0644);
    if (dst_fd < 0) { close(src_fd); return -1; }
    
    char buf[65536];
    ssize_t n;
    while ((n = read(src_fd, buf, sizeof(buf))) > 0) {
        if (writen(dst_fd, buf, n) != n) { close(src_fd); close(dst_fd); return -1; }
    }
    
    close(src_fd);
    close(dst_fd);
    return 0;
}

// ── Simple fast hash for checksum-based skip ──────────────────────────────
static uint64_t fast_hash_file(const char *path, uint64_t file_size) {
    // For small files (< 4KB), hash entire file
    // For large files, hash first 4KB + last 4KB + size
    FILE *f = fopen(path, "rb");
    if (!f) return 0;
    
    uint64_t hash = file_size * 0x9e3779b97f4a7c15ULL;
    
    // Hash first chunk
    unsigned char buf[4096];
    size_t n = fread(buf, 1, sizeof(buf), f);
    for (size_t i = 0; i < n; i++) {
        hash = hash * 0x880355f21e15d8c5ULL + buf[i];
    }
    
    // For files > 8KB, also hash last chunk
    if (file_size > 8192) {
        if (fseeko(f, -(off_t)sizeof(buf), SEEK_END) == 0) {
            n = fread(buf, 1, sizeof(buf), f);
            for (size_t i = 0; i < n; i++) {
                hash = hash * 0x880355f21e15d8c5ULL + buf[i];
            }
        }
    }
    
    fclose(f);
    return hash;
}

// ─── Per-connection context ────────────────────────────────────────────────

typedef struct {
    int tcp_fd;
    int use_udp;
    int udp_fd;
    struct sockaddr_in udp_server_addr;
    uint64_t udp_session_id;
} sync_ctx_t;

// ─── Global stats (atomic for lock-free updates) ──────────────────────────

static atomic_uint_fast64_t g_bytes_sent = 0;
static atomic_uint_fast32_t g_files_sent = 0;

// ─── UDP session setup (called once per worker if -u) ─────────────────────

static int setup_udp(sync_ctx_t *ctx, const char *host) {
    uint32_t req = MAGIC_UDP_REQ;
    if (writen(ctx->tcp_fd, &req, sizeof(req)) != sizeof(req)) return -1;

    uint32_t resp = 0;
    if (readn(ctx->tcp_fd, &resp, sizeof(resp)) != sizeof(resp)) return -1;
    if (resp != RESP_USE_UDP) {
        fprintf(stderr, "[thread] Server refused UDP, using TCP.\n");
        ctx->use_udp = 0;
        return 0;
    }

    udp_handshake_t hs;
    if (readn(ctx->tcp_fd, &hs, sizeof(hs)) != sizeof(hs)) return -1;
    ctx->udp_session_id = hs.session_id;

    ctx->udp_fd = socket(AF_INET, SOCK_DGRAM, 0);
    if (ctx->udp_fd < 0) { perror("UDP socket"); return -1; }

    // Tune UDP socket for better throughput
    int bufsize = 4 * 1024 * 1024; // 4 MB
    setsockopt(ctx->udp_fd, SOL_SOCKET, SO_SNDBUF, &bufsize, sizeof(bufsize));
    setsockopt(ctx->udp_fd, SOL_SOCKET, SO_RCVBUF, &bufsize, sizeof(bufsize));

    memset(&ctx->udp_server_addr, 0, sizeof(ctx->udp_server_addr));
    ctx->udp_server_addr.sin_family = AF_INET;
    ctx->udp_server_addr.sin_port   = htons(hs.udp_port);
    inet_pton(AF_INET, host, &ctx->udp_server_addr.sin_addr);

    // Knock
    udp_packet_header_t knock;
    memset(&knock, 0, sizeof(knock));
    knock.magic      = MAGIC_UDP_KNOCK;
    knock.session_id = hs.session_id;
    sendto(ctx->udp_fd, &knock, sizeof(knock), 0,
           (struct sockaddr *)&ctx->udp_server_addr,
           sizeof(ctx->udp_server_addr));

    uint32_t ack = 0;
    if (readn(ctx->tcp_fd, &ack, sizeof(ack)) != sizeof(ack) || ack != RESP_OK) {
        fprintf(stderr, "[thread] UDP knock failed\n");
        return -1;
    }
    printf("[thread] UDP session ready (port %u)\n", hs.udp_port);
    return 0;
}

// ─── Batch file sending ────────────────────────────────────────────────────

static int send_batch_files(sync_ctx_t *ctx, file_task_t **tasks, int count) {
    // Collect and stat all files first (single pass to avoid double stat)
    // Also pre-open file descriptors for pipelining
    int *fds = malloc(count * sizeof(int));
    struct stat *stats = malloc(count * sizeof(struct stat));
    if (!fds || !stats) { perror("malloc"); free(fds); free(stats); return -1; }

    batch_meta_header_t batch_hdr;
    batch_hdr.magic = MAGIC_BATCH_META;
    batch_hdr.count = count;
    batch_hdr.total_name_len = 0;
    batch_hdr.total_size = 0;

    for (int i = 0; i < count; i++) {
        if (stat(tasks[i]->full_path, &stats[i]) < 0) { 
            perror("stat"); 
            for (int j = 0; j < i; j++) close(fds[j]);
            free(fds); free(stats); 
            return -1; 
        }
        batch_hdr.total_name_len += strlen(tasks[i]->rel_path);
        batch_hdr.total_size += stats[i].st_size;
        
        // Pre-open file descriptors
        fds[i] = open(tasks[i]->full_path, O_RDONLY);
        if (fds[i] < 0) {
            perror("open");
            for (int j = 0; j < i; j++) close(fds[j]);
            free(fds); free(stats);
            return -1;
        }
    }

    // Send batch header
    if (writen(ctx->tcp_fd, &batch_hdr, sizeof(batch_hdr)) != sizeof(batch_hdr)) {
        for (int j = 0; j < count; j++) close(fds[j]);
        free(fds); free(stats);
        return -1;
    }

    // Send files with pipelining: meta+name+data for each file sequentially
    // This allows server to start writing file N while receiving meta for file N+1
    for (int i = 0; i < count; i++) {
        file_meta_t meta;
        meta.magic = 0;
        meta.name_len = strlen(tasks[i]->rel_path);
        meta.file_size = stats[i].st_size;
        meta.mode = stats[i].st_mode;
        meta.checksum = tasks[i]->checksum;

        if (writen(ctx->tcp_fd, &meta, sizeof(meta)) != sizeof(meta)) {
            for (int j = 0; j < count; j++) close(fds[j]);
            free(fds); free(stats);
            return -1;
        }
        if (writen(ctx->tcp_fd, tasks[i]->rel_path, meta.name_len) != (ssize_t)meta.name_len) {
            for (int j = 0; j < count; j++) close(fds[j]);
            free(fds); free(stats);
            return -1;
        }

        // Send file data using sendfile() for zero-copy transfer
        off_t file_offset = 0;
        ssize_t total_sent = 0;
        while (file_offset < (off_t)stats[i].st_size) {
            ssize_t sent = sendfile(ctx->tcp_fd, fds[i], &file_offset, stats[i].st_size - file_offset);
            if (sent < 0) {
                if (errno == EINTR) continue;
                perror("sendfile");
                for (int j = 0; j < count; j++) close(fds[j]);
                free(fds); free(stats);
                return -1;
            }
            if (sent == 0) break; // EOF
            total_sent += sent;
        }
        close(fds[i]);
        atomic_fetch_add(&g_bytes_sent, (uint64_t)total_sent);
        printf("[thread] Sent (pipelined+sendfile): %s\n", tasks[i]->rel_path);
    }

    free(fds);
    free(stats);

    // Update file count atomically
    atomic_fetch_add(&g_files_sent, count);

    return 0;
}

// ─── Worker thread: owns one TCP connection, drains work queue ─────────────

static void *worker_thread(void *arg) {
    worker_arg_t *wa = (worker_arg_t *)arg;

    // Connect TCP
    sync_ctx_t ctx;
    memset(&ctx, 0, sizeof(ctx));
    ctx.use_udp = wa->use_udp;

    ctx.tcp_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (ctx.tcp_fd < 0) { perror("socket"); wa->result = -1; return NULL; }

    // TCP tuning for better throughput
    int optval = 1;
    setsockopt(ctx.tcp_fd, SOL_TCP, TCP_NODELAY, &optval, sizeof(optval));
    setsockopt(ctx.tcp_fd, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(optval));
    
    // Large send/receive buffers
    int bufsize = 4 * 1024 * 1024; // 4 MB
    setsockopt(ctx.tcp_fd, SOL_SOCKET, SO_SNDBUF, &bufsize, sizeof(bufsize));
    setsockopt(ctx.tcp_fd, SOL_SOCKET, SO_RCVBUF, &bufsize, sizeof(bufsize));

    struct sockaddr_in sa;
    memset(&sa, 0, sizeof(sa));
    sa.sin_family = AF_INET;
    sa.sin_port   = htons(wa->port);
    inet_pton(AF_INET, wa->host, &sa.sin_addr);

    if (connect(ctx.tcp_fd, (struct sockaddr *)&sa, sizeof(sa)) < 0) {
        perror("connect"); close(ctx.tcp_fd); wa->result = -1; return NULL;
    }
    printf("[thread %d] TCP connected to %s:%d\n", wa->thread_id, wa->host, wa->port);

    if (wa->use_udp && setup_udp(&ctx, wa->host) < 0) {
        close(ctx.tcp_fd); wa->result = -1; return NULL;
    }

    // Drain queue with batching
    file_task_t *batch[BATCH_MAX_FILES];
    int batch_count = 0;

    file_task_t *task;
    while ((task = wq_pop(wa->queue)) != NULL) {
        batch[batch_count++] = task;

        // If batch is full, send it
        if (batch_count >= BATCH_MAX_FILES) {
            if (send_batch_files(&ctx, batch, batch_count) < 0) {
                fprintf(stderr, "[thread %d] Batch send failed\n", wa->thread_id);
                wa->result = -1;
            }
            // Free all tasks in batch
            for (int i = 0; i < batch_count; i++) {
                free(batch[i]);
            }
            batch_count = 0;
        }
    }

    // Send remaining files in partial batch
    if (batch_count > 0) {
        if (send_batch_files(&ctx, batch, batch_count) < 0) {
            fprintf(stderr, "[thread %d] Final batch send failed\n", wa->thread_id);
            wa->result = -1;
        }
        for (int i = 0; i < batch_count; i++) {
            free(batch[i]);
        }
    }

    // Signal done to server
    file_meta_t done;
    memset(&done, 0, sizeof(done));
    done.magic = MAGIC_DONE;
    writen(ctx.tcp_fd, &done, sizeof(done));

    if (ctx.use_udp && ctx.udp_fd > 0) close(ctx.udp_fd);
    close(ctx.tcp_fd);
    printf("[thread %d] Done.\n", wa->thread_id);
    return NULL;
}

// ─── Scanner: recursively walks dir, collects files, sorts by size, pushes to queue ──

typedef struct {
    work_queue_t *queue;
    char          base_path[1024];
} scanner_arg_t;

// Comparison function for sorting files by size (ascending: small files first)
static int compare_files_by_size(const void *a, const void *b) {
    const file_task_t *fa = *(const file_task_t **)a;
    const file_task_t *fb = *(const file_task_t **)b;
    if (fa->file_size < fb->file_size) return -1;
    if (fa->file_size > fb->file_size) return 1;
    return 0;
}

// Collect all files recursively into a list
static int collect_files(const char *base, const char *sub, file_task_t **list, int *count, int max_files) {
    char cur[2048];
    if (sub && sub[0])
        snprintf(cur, sizeof(cur), "%s/%s", base, sub);
    else
        snprintf(cur, sizeof(cur), "%s", base);

    DIR *dir = opendir(cur);
    if (!dir) { perror("opendir"); return -1; }

    struct dirent *e;
    while ((e = readdir(dir)) != NULL) {
        if (!strcmp(e->d_name, ".") || !strcmp(e->d_name, "..")) continue;

        char rel[1024], full[2048];
        if (sub && sub[0]) {
            snprintf(rel,  sizeof(rel),  "%s/%s", sub, e->d_name);
        } else {
            snprintf(rel,  sizeof(rel),  "%s", e->d_name);
        }
        snprintf(full, sizeof(full), "%s/%s", base, rel);
        full[sizeof(full) - 1] = '\0';
        rel[sizeof(rel) - 1] = '\0';

        struct stat st;
        if (lstat(full, &st) < 0) { perror("lstat"); continue; }

        if (S_ISDIR(st.st_mode)) {
            if (collect_files(base, rel, list, count, max_files) < 0) {
                closedir(dir);
                return -1;
            }
        } else if (S_ISREG(st.st_mode)) {
            if (*count >= max_files) {
                fprintf(stderr, "Warning: Maximum file count reached, stopping scan\n");
                closedir(dir);
                return 0;
            }
            list[*count] = malloc(sizeof(file_task_t));
            if (!list[*count]) { perror("malloc"); closedir(dir); return -1; }
            snprintf(list[*count]->full_path, sizeof(list[*count]->full_path), "%s", full);
            snprintf(list[*count]->rel_path, sizeof(list[*count]->rel_path), "%s", rel);
            list[*count]->file_size = st.st_size;
            list[*count]->checksum = fast_hash_file(full, st.st_size);
            list[*count]->next = NULL;
            (*count)++;
        }
    }
    closedir(dir);
    return 0;
}

static void *scanner_thread(void *arg) {
    scanner_arg_t *sa = (scanner_arg_t *)arg;
    
    // First pass: count files to determine array size
    int file_count = 0;
    char cur[2048];
    DIR *dir = opendir(sa->base_path);
    if (!dir) { perror("opendir"); wq_finish(sa->queue); return NULL; }
    
    struct dirent *e;
    while ((e = readdir(dir)) != NULL) {
        if (!strcmp(e->d_name, ".") || !strcmp(e->d_name, "..")) continue;
        
        snprintf(cur, sizeof(cur), "%s/%s", sa->base_path, e->d_name);
        struct stat st;
        if (lstat(cur, &st) < 0) continue;
        
        if (S_ISDIR(st.st_mode)) {
            // Recursively count
            DIR *subdir = opendir(cur);
            if (subdir) {
                struct dirent *se;
                while ((se = readdir(subdir)) != NULL) {
                    if (!strcmp(se->d_name, ".") || !strcmp(se->d_name, "..")) continue;
                    file_count++;
                }
                closedir(subdir);
            }
        } else if (S_ISREG(st.st_mode)) {
            file_count++;
        }
    }
    closedir(dir);
    
    // Allocate array for file tasks
    file_task_t **file_list = malloc(file_count * sizeof(file_task_t *));
    if (!file_list) { perror("malloc"); wq_finish(sa->queue); return NULL; }
    
    int count = 0;
    if (collect_files(sa->base_path, "", file_list, &count, file_count) < 0) {
        for (int i = 0; i < count; i++) free(file_list[i]);
        free(file_list);
        wq_finish(sa->queue);
        return NULL;
    }
    
    // Sort files by size (smallest first for better pipelining)
    qsort(file_list, count, sizeof(file_task_t *), compare_files_by_size);
    
    // Push sorted files to queue
    for (int i = 0; i < count; i++) {
        wq_push(sa->queue, file_list[i]->full_path, file_list[i]->rel_path, file_list[i]->file_size, file_list[i]->checksum);
        free(file_list[i]);
    }
    free(file_list);
    
    wq_finish(sa->queue);   /* broadcast to all waiting workers */
    return NULL;
}

// ─── main ──────────────────────────────────────────────────────────────────

static void print_usage(const char *prog) {
    fprintf(stderr,
        "Usage: %s -h <host> [-p <port>] -s <source> [-n <connections>] [-u]\n"
        "  -n  number of parallel TCP connections (default: 4)\n"
        "  -u  enable UDP data channel\n", prog);
}

int main(int argc, char *argv[]) {
    char *host    = NULL;
    int   port    = DEFAULT_PORT;
    char *source  = NULL;
    int   nconn   = 4;
    int   use_udp = 0;
    int   opt;

    while ((opt = getopt(argc, argv, "h:p:s:n:u")) != -1) {
        switch (opt) {
            case 'h': host    = optarg;       break;
            case 'p': port    = atoi(optarg); break;
            case 's': source  = optarg;       break;
            case 'n': nconn   = atoi(optarg); break;
            case 'u': use_udp = 1;            break;
            default:  print_usage(argv[0]);   exit(EXIT_FAILURE);
        }
    }

    if (!host || !source) { print_usage(argv[0]); exit(EXIT_FAILURE); }
    if (nconn < 1) nconn = 1;
    if (nconn > 16) nconn = 16;

    struct stat st;
    if (stat(source, &st) < 0) { perror("stat source"); exit(EXIT_FAILURE); }

    printf("Starting sync: %s → %s:%d  (connections=%d, udp=%s)\n",
           source, host, port, nconn, use_udp ? "yes" : "no");

    // ── Create temp dir for single files (to use directory scanning) ─────────
    char temp_dir[2048] = "";
    int use_temp_dir = 0;
    char *final_source = (char *)source;
    
    if (S_ISREG(st.st_mode)) {
        // Create temp dir, copy file there with its name
        strncpy(temp_dir, "/tmp/fastsync_temp_XXXXXX", sizeof(temp_dir));
        if (mkdtemp(temp_dir) == NULL) {
            perror("mkdtemp");
            exit(EXIT_FAILURE);
        }
        use_temp_dir = 1;
        
        char *fname = strrchr(source, '/');
        char new_path[2048];
        snprintf(new_path, sizeof(new_path), "%s/%s", temp_dir, fname ? fname + 1 : source);
        
        // Copy file to temp dir
        if (copy_file(source, new_path) < 0) {
            perror("copy_file");
            rmdir(temp_dir);
            exit(EXIT_FAILURE);
        }
        final_source = temp_dir;
    }

    // ── Init shared work queue ────────────────────────────────────────────
    work_queue_t queue;
    wq_init(&queue);

    // ── Start scanner thread ──────────────────────────────────────────────
    scanner_arg_t sarg;
    sarg.queue = &queue;
    snprintf(sarg.base_path, sizeof(sarg.base_path), "%s", final_source);
    pthread_t stid;
    pthread_create(&stid, NULL, scanner_thread, &sarg);
    pthread_detach(stid);

    // ── Spawn N worker threads ────────────────────────────────────────────
    pthread_t   *tids = malloc(nconn * sizeof(pthread_t));
    worker_arg_t *args = malloc(nconn * sizeof(worker_arg_t));
    if (!tids || !args) { perror("malloc"); exit(EXIT_FAILURE); }

    // ── Start timer ───────────────────────────────────────────────────────
    struct timespec t_start, t_end;
    clock_gettime(CLOCK_MONOTONIC, &t_start);

    for (int i = 0; i < nconn; i++) {
        args[i].thread_id = i;
        args[i].queue     = &queue;
        args[i].host      = host;
        args[i].port      = port;
        args[i].use_udp   = use_udp;
        args[i].result    = 0;
        pthread_create(&tids[i], NULL, worker_thread, &args[i]);
    }

    // ── Wait for all workers ──────────────────────────────────────────────
    int overall = 0;
    for (int i = 0; i < nconn; i++) {
        pthread_join(tids[i], NULL);
        if (args[i].result != 0) overall = -1;
    }

    // ── Stop timer & print summary ────────────────────────────────────────
    clock_gettime(CLOCK_MONOTONIC, &t_end);
    double elapsed_ms = (t_end.tv_sec - t_start.tv_sec) * 1000.0
                      + (t_end.tv_nsec - t_start.tv_nsec) / 1e6;

    uint64_t total_bytes = atomic_load(&g_bytes_sent);
    uint32_t total_files = atomic_load(&g_files_sent);
    double   throughput  = (total_bytes / 1048576.0) / (elapsed_ms / 1000.0);

    printf("\n=== Sync complete: %u files, %lu bytes, %.1f ms, %.2f MB/s ===\n",
           total_files, (unsigned long)total_bytes, elapsed_ms, throughput);
    /* Machine-readable line for benchmark.sh to parse */
    printf("BENCH: connections=%d bytes=%lu ms=%.1f throughput_mbs=%.2f\n",
           nconn, (unsigned long)total_bytes, elapsed_ms, throughput);

    wq_destroy(&queue);
    free(tids);
    free(args);
    
    // Clean up temp directory for single file
    if (use_temp_dir) {
        // Remove the temp directory and its contents
        char cmd[2048];
        snprintf(cmd, sizeof(cmd), "rm -rf %s", temp_dir);
        system(cmd);
    }
    
    return overall == 0 ? EXIT_SUCCESS : EXIT_FAILURE;
}