#include "common.h"
#include <getopt.h>
#include <poll.h>
#include <time.h>
#include <netinet/tcp.h>
#include <stdio.h>

// ── Simple fast hash for checksum-based skip ──────────────────────────────
static uint64_t fast_hash_file(const char *path, uint64_t file_size) {
    FILE *f = fopen(path, "rb");
    if (!f) return 0;
    
    uint64_t hash = file_size * 0x9e3779b97f4a7c15ULL;
    
    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];
    }
    
    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;
}

typedef struct {
    uint32_t file_id;
    char target_path[2048];
    int out_fd;
    uint64_t file_size;
    uint32_t total_blocks;
    char *received_blocks;
} active_file_t;

int setup_udp_session(int tcp_fd, int *udp_fd_out, struct sockaddr_in *client_udp_addr_out, uint64_t *session_id_out) {
    int udp_fd = socket(AF_INET, SOCK_DGRAM, 0);
    if (udp_fd < 0) {
        perror("UDP socket creation failed");
        return -1;
    }

    struct sockaddr_in udp_addr;
    memset(&udp_addr, 0, sizeof(udp_addr));
    udp_addr.sin_family = AF_INET;
    udp_addr.sin_addr.s_addr = INADDR_ANY;
    udp_addr.sin_port = 0; // Ephemeral port

    if (bind(udp_fd, (struct sockaddr *)&udp_addr, sizeof(udp_addr)) < 0) {
        perror("UDP bind failed");
        close(udp_fd);
        return -1;
    }

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

    socklen_t addr_len = sizeof(udp_addr);
    if (getsockname(udp_fd, (struct sockaddr *)&udp_addr, &addr_len) < 0) {
        perror("getsockname failed");
        close(udp_fd);
        return -1;
    }

    uint32_t udp_port = ntohs(udp_addr.sin_port);
    uint64_t session_id = ((uint64_t)rand() << 32) | rand();

    // Send response code and handshake over TCP
    uint32_t resp = RESP_USE_UDP;
    if (writen(tcp_fd, &resp, sizeof(resp)) != sizeof(resp)) {
        close(udp_fd);
        return -1;
    }

    udp_handshake_t handshake;
    handshake.udp_port = udp_port;
    handshake.session_id = session_id;
    if (writen(tcp_fd, &handshake, sizeof(handshake)) != sizeof(handshake)) {
        close(udp_fd);
        return -1;
    }

    printf("UDP port bound to %u. Waiting for knock with Session ID: %lu...\n", udp_port, (unsigned long)session_id);

    // Timeout of 30s for the UDP knock (more reasonable for variable networks)
    struct timeval tv;
    tv.tv_sec = 30;
    tv.tv_usec = 0;
    setsockopt(udp_fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));

    struct sockaddr_in client_udp_addr;
    socklen_t client_len = sizeof(client_udp_addr);
    udp_packet_header_t knock_hdr;

    ssize_t n = recvfrom(udp_fd, &knock_hdr, sizeof(knock_hdr), 0,
                         (struct sockaddr *)&client_udp_addr, &client_len);
    if (n < 0) {
        perror("UDP knock timeout or recvfrom failed");
        close(udp_fd);
        return -1;
    }

    if (knock_hdr.magic != MAGIC_UDP_KNOCK || knock_hdr.session_id != session_id) {
        fprintf(stderr, "Invalid UDP knock magic or session ID\n");
        close(udp_fd);
        return -1;
    }

    // Confirm session established on TCP
    resp = RESP_OK;
    if (writen(tcp_fd, &resp, sizeof(resp)) != sizeof(resp)) {
        close(udp_fd);
        return -1;
    }

    printf("UDP Session established successfully.\n");
    *udp_fd_out = udp_fd;
    *client_udp_addr_out = client_udp_addr;
    *session_id_out = session_id;
    return 0;
}

int receive_file_udp_loop(int tcp_fd, int udp_fd, uint64_t session_id, active_file_t *active) {
    struct pollfd fds[2];
    fds[0].fd = tcp_fd;
    fds[0].events = POLLIN;
    fds[1].fd = udp_fd;
    fds[1].events = POLLIN;

    // Reset SO_RCVTIMEO on UDP socket (block indefinitely during active sync)
    struct timeval tv;
    tv.tv_sec = 0;
    tv.tv_usec = 0;
    setsockopt(udp_fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));

    while (1) {
        int poll_ret = poll(fds, 2, -1);
        if (poll_ret < 0) {
            if (errno == EINTR) {
                continue;
            }
            perror("poll failed");
            return -1;
        }

        // 1. Check UDP Data socket
        if (fds[1].revents & POLLIN) {
            char pkt[sizeof(udp_packet_header_t) + UDP_PAYLOAD_MAX];
            struct sockaddr_in src_addr;
            socklen_t src_len = sizeof(src_addr);
            ssize_t n = recvfrom(udp_fd, pkt, sizeof(pkt), 0,
                                 (struct sockaddr *)&src_addr, &src_len);
            if (n > 0) {
                if (n >= (ssize_t)sizeof(udp_packet_header_t)) {
                    udp_packet_header_t *hdr = (udp_packet_header_t *)pkt;
                    if (hdr->magic == MAGIC_UDP_DATA && hdr->session_id == session_id && hdr->file_id == active->file_id) {
                        uint32_t bid = hdr->block_id;
                        if (bid < active->total_blocks) {
                            if (active->received_blocks[bid] == 0) {
                                char *payload = pkt + sizeof(udp_packet_header_t);
                                if (pwrite(active->out_fd, payload, hdr->payload_len, hdr->block_offset) != (ssize_t)hdr->payload_len) {
                                    perror("pwrite file block failed");
                                } else {
                                    active->received_blocks[bid] = 1;
                                }
                            }
                        }
                    }
                }
            }
        }

        // 2. Check TCP command socket
        if (fds[0].revents & POLLIN) {
            uint32_t req_magic = 0;
            ssize_t n = readn(tcp_fd, &req_magic, sizeof(req_magic));
            if (n <= 0) {
                return -1; // Disconnected
            }

            if (req_magic == MAGIC_VERIFY) {
                uint32_t file_id = 0;
                if (readn(tcp_fd, &file_id, sizeof(file_id)) != sizeof(file_id)) {
                    return -1;
                }

                if (file_id != active->file_id) {
                    fprintf(stderr, "File ID mismatch on verify query: got %u, active is %u\n", file_id, active->file_id);
                    return -1;
                }

                // Count missing blocks
                uint32_t missing_count = 0;
                for (uint32_t i = 0; i < active->total_blocks; i++) {
                    if (active->received_blocks[i] == 0) {
                        missing_count++;
                    }
                }

                // Send count
                if (writen(tcp_fd, &missing_count, sizeof(missing_count)) != sizeof(missing_count)) {
                    return -1;
                }

                if (missing_count == 0) {
                    // Success!
                    break;
                }

                // Compile and send missing block IDs
                uint32_t *missing_array = malloc(missing_count * sizeof(uint32_t));
                if (!missing_array) {
                    perror("malloc missing array");
                    return -1;
                }
                uint32_t idx = 0;
                for (uint32_t i = 0; i < active->total_blocks; i++) {
                    if (active->received_blocks[i] == 0) {
                        missing_array[idx++] = i;
                    }
                }

                if (writen(tcp_fd, missing_array, missing_count * sizeof(uint32_t)) != (ssize_t)(missing_count * sizeof(uint32_t))) {
                    free(missing_array);
                    return -1;
                }
                free(missing_array);
            } else {
                fprintf(stderr, "Received unexpected TCP command magic: 0x%08x\n", req_magic);
                return -1;
            }
        }
    }

    return 0;
}

// ─── Per-connection thread wrapper ───────────────────────────────────────────

/* Forward declaration — defined further below */
void handle_client(int sock_fd, const char *dest_dir);

typedef struct {

    int   client_fd;
    char  dest_dir[1024];
} client_thread_arg_t;

static void *client_thread(void *arg) {
    client_thread_arg_t *ca = (client_thread_arg_t *)arg;
    handle_client(ca->client_fd, ca->dest_dir);
    close(ca->client_fd);
    free(ca);
    return NULL;
}

void handle_client(int sock_fd, const char *dest_dir) {
    printf("Client connected. Starting transfer session...\n");

    int use_udp = 0;
    int udp_fd = -1;
    struct sockaddr_in client_udp_addr;
    uint64_t udp_session_id = 0;
    uint32_t file_id_counter = 0;

    while (1) {
        uint32_t magic = 0;
        // Peek or read magic
        ssize_t n = readn(sock_fd, &magic, sizeof(magic));
        if (n <= 0) {
            if (n < 0) {
                perror("read magic failed");
            } else {
                printf("Client disconnected.\n");
            }
            break;
        }

        if (magic == MAGIC_UDP_REQ) {
            if (setup_udp_session(sock_fd, &udp_fd, &client_udp_addr, &udp_session_id) < 0) {
                fprintf(stderr, "Failed to initialize UDP session. Falling back to TCP.\n");
                uint32_t resp_fail = RESP_ERROR;
                writen(sock_fd, &resp_fail, sizeof(resp_fail));
            } else {
                use_udp = 1;
            }
            continue;
        }

        if (magic == MAGIC_DONE) {
            printf("Sync completed successfully.\n");
            break;
        }

        if (magic == MAGIC_BATCH_META) {
            // Handle batch metadata
            // We already read the magic (4 bytes), now read the rest of the header (16 bytes)
            batch_meta_header_t batch_hdr;
            batch_hdr.magic = magic;
            if (readn(sock_fd, &batch_hdr.count, sizeof(batch_hdr.count)) != sizeof(batch_hdr.count) ||
                readn(sock_fd, &batch_hdr.total_name_len, sizeof(batch_hdr.total_name_len)) != sizeof(batch_hdr.total_name_len) ||
                readn(sock_fd, &batch_hdr.total_size, sizeof(batch_hdr.total_size)) != sizeof(batch_hdr.total_size)) {
                fprintf(stderr, "Failed to read batch header\n");
                break;
            }

            if (batch_hdr.count > BATCH_MAX_FILES) {
                fprintf(stderr, "Batch count %u exceeds max %d\n", batch_hdr.count, BATCH_MAX_FILES);
                break;
            }

            // No response needed - client sends data immediately after each file's metadata
            // Process each file in batch with pipelining: meta+name+data for each file
            uint32_t first_file_id = file_id_counter + 1;

            for (uint32_t i = 0; i < batch_hdr.count; i++) {
                // Read metadata for this file
                file_meta_t meta;
                if (readn(sock_fd, &meta, sizeof(meta)) != sizeof(meta)) {
                    fprintf(stderr, "Failed to read pipelined file meta %u\n", i);
                    break;
                }

                // Read filename
                char *filename = malloc(meta.name_len + 1);
                if (!filename) {
                    perror("malloc filename");
                    break;
                }
                if (readn(sock_fd, filename, meta.name_len) != (ssize_t)meta.name_len) {
                    fprintf(stderr, "Failed to read pipelined filename %u\n", i);
                    free(filename);
                    break;
                }
                filename[meta.name_len] = '\0';

                char target_path[2048];
                snprintf(target_path, sizeof(target_path), "%s/%s", dest_dir, filename);

                // Check if file already exists with same size and checksum (skip if so)
                struct stat existing_st;
                if (stat(target_path, &existing_st) == 0 && 
                    (uint64_t)existing_st.st_size == meta.file_size &&
                    S_ISREG(existing_st.st_mode)) {
                    // File exists with same size - check checksum
                    uint64_t existing_checksum = fast_hash_file(target_path, meta.file_size);
                    if (existing_checksum == meta.checksum) {
                        printf("Skip (checksum match): %s\n", target_path);
                        free(filename);
                        continue;
                    }
                }

                // Create parent directory
                char *slash = strrchr(target_path, '/');
                if (slash) {
                    *slash = '\0';
                    if (make_path(target_path, 0755) < 0) {
                        perror("Failed to create parent directory path");
                        *slash = '/';
                        free(filename);
                        continue;
                    }
                    *slash = '/';
                }

                int out_fd = open(target_path, O_WRONLY | O_CREAT | O_TRUNC, meta.mode);
                if (out_fd < 0) {
                    perror("Failed to open target file");
                    free(filename);
                    continue;
                }

                printf("Receiving (pipelined): %s (%lu bytes)\n", filename, (unsigned long)meta.file_size);

                // Read file data
                char buffer[65536];
                uint64_t bytes_left = meta.file_size;
                int write_error = 0;
                while (bytes_left > 0) {
                    size_t to_read = (bytes_left > sizeof(buffer)) ? sizeof(buffer) : bytes_left;
                    ssize_t nread = read(sock_fd, buffer, to_read);
                    if (nread < 0) {
                        if (errno == EINTR) continue;
                        perror("Socket read error");
                        write_error = 1;
                        break;
                    }
                    if (nread == 0) {
                        fprintf(stderr, "Unexpected socket EOF\n");
                        write_error = 1;
                        break;
                    }
                    if (!write_error) {
                        if (writen(out_fd, buffer, nread) != nread) {
                            perror("File write error");
                            write_error = 1;
                        }
                    }
                    bytes_left -= nread;
                }
                if (!write_error) {
                    fsync(out_fd);  // Ensure data is flushed to disk
                }
                close(out_fd);
                free(filename);

                if (!write_error) {
                    printf("Saved (pipelined TCP): %s\n", target_path);
                }
            }

            file_id_counter = first_file_id + batch_hdr.count;
            continue;
        }

        if (magic != MAGIC_META) {
            fprintf(stderr, "Received invalid magic header: 0x%08x\n", magic);
            break;
        }

        // We received a file metadata packet (read rest of it)
        uint32_t name_len = 0;
        uint64_t file_size = 0;
        uint32_t mode = 0;

        if (readn(sock_fd, &name_len, sizeof(name_len)) != sizeof(name_len) ||
            readn(sock_fd, &file_size, sizeof(file_size)) != sizeof(file_size) ||
            readn(sock_fd, &mode, sizeof(mode)) != sizeof(mode)) {
            fprintf(stderr, "Failed to read file metadata fields\n");
            break;
        }

        // Read filename
        char filename[1024];
        if (name_len >= sizeof(filename)) {
            fprintf(stderr, "Filename too long: %u bytes\n", name_len);
            break;
        }
        if (readn(sock_fd, filename, name_len) != (ssize_t)name_len) {
            fprintf(stderr, "Failed to read filename\n");
            break;
        }
        filename[name_len] = '\0';

        char target_path[2048];
        snprintf(target_path, sizeof(target_path), "%s/%s", dest_dir, filename);

        // Ensure directories exist
        if (make_path(target_path, 0755) < 0) {
            perror("Failed to create parent directory path");
            file_response_t resp_err = {RESP_ERROR, 0, {0, 0}};
            writen(sock_fd, &resp_err, sizeof(resp_err));
            continue;
        }

        uint32_t file_id = ++file_id_counter;

        if (use_udp) {
            file_response_t resp_udp = {RESP_USE_UDP, file_id, {0, 0}};
            if (writen(sock_fd, &resp_udp, sizeof(resp_udp)) != sizeof(resp_udp)) {
                perror("Failed to send file response header");
                break;
            }

            active_file_t active;
            active.file_id = file_id;
            active.file_size = file_size;
            active.total_blocks = (file_size + UDP_PAYLOAD_MAX - 1) / UDP_PAYLOAD_MAX;
            active.received_blocks = calloc(active.total_blocks, 1);
            snprintf(active.target_path, sizeof(active.target_path), "%s", target_path);

            active.out_fd = open(target_path, O_WRONLY | O_CREAT | O_TRUNC, mode);
            if (active.out_fd < 0) {
                perror("Failed to open target file for UDP sync");
                free(active.received_blocks);
                break;
            }

            printf("Receiving over UDP: %s (%lu bytes, %u blocks)\n", filename, (unsigned long)file_size, active.total_blocks);

            int loop_ret = receive_file_udp_loop(sock_fd, udp_fd, udp_session_id, &active);
            close(active.out_fd);
            free(active.received_blocks);

            if (loop_ret < 0) {
                fprintf(stderr, "UDP receive loop failed\n");
                break;
            }

            // Send final ACK over TCP
            uint32_t ack_ok = RESP_OK;
            if (writen(sock_fd, &ack_ok, sizeof(ack_ok)) != sizeof(ack_ok)) {
                perror("Failed to send final file ACK");
                break;
            }

            printf("Saved (UDP): %s\n", target_path);
        } else {
            // TCP fallback
            file_response_t resp_tcp = {RESP_SEND_DATA, 0, {0, 0}};
            if (writen(sock_fd, &resp_tcp, sizeof(resp_tcp)) != sizeof(resp_tcp)) {
                perror("Failed to send file response header");
                break;
            }

            int out_fd = open(target_path, O_WRONLY | O_CREAT | O_TRUNC, mode);
            if (out_fd < 0) {
                perror("Failed to open target file for TCP sync");
                break;
            }

            printf("Receiving over TCP: %s (%lu bytes)\n", filename, (unsigned long)file_size);

            char buffer[65536];
            uint64_t bytes_left = file_size;
            int write_error = 0;
            while (bytes_left > 0) {
                size_t to_read = (bytes_left > sizeof(buffer)) ? sizeof(buffer) : bytes_left;
                ssize_t nread = read(sock_fd, buffer, to_read);
                if (nread < 0) {
                    if (errno == EINTR) {
                        continue;
                    }
                    perror("Socket read error");
                    write_error = 1;
                    break;
                }
                if (nread == 0) {
                    fprintf(stderr, "Unexpected socket EOF\n");
                    write_error = 1;
                    break;
                }

                if (!write_error) {
                    if (writen(out_fd, buffer, nread) != nread) {
                        perror("File write error");
                        write_error = 1;
                    }
                }
                bytes_left -= nread;
            }
            close(out_fd);

            uint32_t ack = write_error ? RESP_ERROR : RESP_OK;
            if (writen(sock_fd, &ack, sizeof(ack)) != sizeof(ack)) {
                perror("Failed to send final file ACK");
                break;
            }

            if (!write_error) {
                printf("Saved (TCP): %s\n", target_path);
            }
        }
    }

    if (udp_fd >= 0) {
        close(udp_fd);
    }
}

void print_usage(const char *prog) {
    fprintf(stderr, "Usage: %s [-p <port>] -d <destination_dir>\n", prog);
}

int main(int argc, char *argv[]) {
    srand(time(NULL));
    int port = DEFAULT_PORT;
    char *dest_dir = NULL;
    int opt;

    while ((opt = getopt(argc, argv, "p:d:")) != -1) {
        switch (opt) {
            case 'p':
                port = atoi(optarg);
                break;
            case 'd':
                dest_dir = optarg;
                break;
            default:
                print_usage(argv[0]);
                exit(EXIT_FAILURE);
        }
    }

    if (!dest_dir) {
        print_usage(argv[0]);
        exit(EXIT_FAILURE);
    }

    struct stat st;
    if (stat(dest_dir, &st) < 0) {
        if (mkdir(dest_dir, 0755) < 0) {
            perror("mkdir dest_dir");
            exit(EXIT_FAILURE);
        }
    } else if (!S_ISDIR(st.st_mode)) {
        fprintf(stderr, "Destination path '%s' is not a directory.\n", dest_dir);
        exit(EXIT_FAILURE);
    }

    int server_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (server_fd < 0) {
        perror("socket creation failed");
        exit(EXIT_FAILURE);
    }

    int optval = 1;
    if (setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(optval)) < 0) {
        perror("setsockopt SO_REUSEADDR failed");
        close(server_fd);
        exit(EXIT_FAILURE);
    }

    struct sockaddr_in address;
    memset(&address, 0, sizeof(address));
    address.sin_family = AF_INET;
    address.sin_addr.s_addr = INADDR_ANY;
    address.sin_port = htons(port);

    if (bind(server_fd, (struct sockaddr *)&address, sizeof(address)) < 0) {
        perror("bind failed");
        close(server_fd);
        exit(EXIT_FAILURE);
    }

    if (listen(server_fd, 32) < 0) {
        perror("listen failed");
        close(server_fd);
        exit(EXIT_FAILURE);
    }

    printf("Server listening on port %d, writing to directory: %s\n", port, dest_dir);

    while (1) {
        struct sockaddr_in client_addr;
        socklen_t addr_len = sizeof(client_addr);
        int client_fd = accept(server_fd, (struct sockaddr *)&client_addr, &addr_len);
        if (client_fd < 0) {
            if (errno == EINTR) {
                continue;
            }
            perror("accept failed");
            continue;
        }

        char client_ip[INET_ADDRSTRLEN];
        inet_ntop(AF_INET, &client_addr.sin_addr, client_ip, sizeof(client_ip));
        printf("\nConnection accepted from %s:%d\n", client_ip, ntohs(client_addr.sin_port));

        client_thread_arg_t *ca = malloc(sizeof(client_thread_arg_t));
        if (!ca) { perror("malloc"); close(client_fd); continue; }
        ca->client_fd = client_fd;
        snprintf(ca->dest_dir, sizeof(ca->dest_dir), "%s", dest_dir);

        // TCP tuning for the accepted socket
        int optval = 1;
        setsockopt(client_fd, SOL_TCP, TCP_NODELAY, &optval, sizeof(optval));
        int bufsize = 4 * 1024 * 1024; // 4 MB
        setsockopt(client_fd, SOL_SOCKET, SO_SNDBUF, &bufsize, sizeof(bufsize));
        setsockopt(client_fd, SOL_SOCKET, SO_RCVBUF, &bufsize, sizeof(bufsize));

        pthread_t tid;
        if (pthread_create(&tid, NULL, client_thread, ca) != 0) {
            perror("pthread_create");
            close(client_fd);
            free(ca);
        } else {
            pthread_detach(tid);  /* fire-and-forget; thread cleans up itself */
        }
    }

    close(server_fd);
    return EXIT_SUCCESS;
}
