#include <stdlib.h>
#include <stdio.h>
#include <netdb.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <string.h>
#include <math.h>

#ifdef __APPLE__
#include "./endian.h"
#else
#include <endian.h>
#endif

#include "./cascade.h"
#include "./sha256.h"

char tracker_ip[IP_LEN];
char tracker_port[PORT_LEN];
char my_ip[IP_LEN];
char my_port[PORT_LEN];

struct csc_file *casc_file;
csc_block_t** queue;
csc_peer_t* peers;

void free_resources()
{
    free(queue);
    free(peers);
    csc_free_file(casc_file);
}

unsigned char* get_file_sha(const char* sourcefile, uint8_t hash[32], int size)
{
    int casc_file_size;

    FILE* fp = fopen(sourcefile, "rb");
    if (fp == 0)
    {
        printf("Failed to open source: %s\n", sourcefile);
        return NULL;
    }

    fseek(fp, 0L, SEEK_END);
    casc_file_size = ftell(fp);
    fseek(fp, 0L, SEEK_SET);

    char buffer[casc_file_size];
    fread(buffer, casc_file_size, 1, fp);
    fclose(fp);

    SHA256_CTX shactx;
    uint8_t shabuffer[size];
    sha256_init(&shactx);
    sha256_update(&shactx, buffer, casc_file_size);
    sha256_final(&shactx, &shabuffer);

    for (int i=0; i<size; i++)
    {
        hash[i] = shabuffer[i];
    }
}

void download_only_peer(char *cascade_file)
{
    printf("Managing download only for: %s\n", cascade_file);
    if (access(cascade_file, F_OK ) != 0 )
    {
        fprintf(stderr, ">> File %s does not exist\n", cascade_file);
        exit(EXIT_FAILURE);
    }

    char output_file[strlen(cascade_file)];
    memcpy(output_file, cascade_file, strlen(cascade_file));
    char* r = strstr(cascade_file, "cascade");
    int cutoff = r - cascade_file ;
    output_file[cutoff-1] = '\0';
    printf("Downloading to: %s\n", output_file);

    casc_file = csc_parse_file(cascade_file, output_file);
    csc_block_t** uncomp = malloc(sizeof(csc_block_t*) * casc_file->blockcount);
    int uncomp_count = 0;
    for (unsigned long long i = 0;i<casc_file->blockcount;i++) {
        if ((&casc_file->blocks[i])->completed == 0) {
            uncomp[uncomp_count] = &casc_file->blocks[i];
            uncomp_count++;
        }
    }

    queue = malloc(sizeof(csc_block_t*) * uncomp_count);
    for (int i = 0;i<uncomp_count;i++) {
        queue[i] = uncomp[i];
    }
    free(uncomp);

    uint8_t hash_buf[32];
    get_file_sha(cascade_file, hash_buf, 32);

    int peercount = 0;
    while (peercount == 0)
    {
        peercount = get_peers_list(&peers, hash_buf);
        if (peercount == 0)
        {
            printf("No peers were found. Will try again in %d seconds\n", PEER_REQUEST_DELAY);
            fflush(stdout);
            sleep(PEER_REQUEST_DELAY);
        }
        else
        {
            printf("Found %d peer(s)\n", peercount);
        }
    }

    csc_peer_t peer = (peers[0]);
    // Get a good peer if one is available
    for (int i=0; i<peercount; i++)
    {
        if (peers[i].good)
        {
            peer = peers[i];
        }
    }

    for (int i=0; i<uncomp_count; i++)
    {
        get_block(queue[i], peer, hash_buf, output_file);
    }

    free_resources();
}

int count_occurences(char string[], char c)
{
    int i=0;
    int count=0;
    for(i=0; i<strlen(string); i++)
    {
        if(string[i] == c)
        {
            count++;
        }
    }
    return count;
}

// Adapted from: https://stackoverflow.com/a/35452093/
uint8_t* hex_to_bytes(const char* string) {

    if(string == NULL)
       return NULL;

    size_t slength = strlen(string);
    if((slength % 2) != 0) // must be even
       return NULL;

    size_t dlength = slength / 2;

    uint8_t* data = malloc(dlength);
    memset(data, 0, dlength);

    size_t index = 0;
    while (index < slength) {
        char c = string[index];
        int value = 0;
        if(c >= '0' && c <= '9')
          value = (c - '0');
        else if (c >= 'A' && c <= 'F')
          value = (10 + (c - 'A'));
        else if (c >= 'a' && c <= 'f')
          value = (10 + (c - 'a'));
        else {
          free(data);
          return NULL;
        }

        data[(index/2)] += value << (((index + 1) % 2) * 4);

        index++;
    }

    return data;
}

/*
 * Parses a cascade file, given the sourcepath input and destination, which may or may not exist.
 * Returns a pointer to a datastructure describing the file, or NULL if the file could not be parsed
 */
csc_file_t* csc_parse_file(const char* sourcefile, const char* destination)
{
    FILE* fp = fopen(sourcefile, "rb");
    if (fp == 0)
    {
        printf("Failed to open source: %s\n", sourcefile);
        return NULL;
    }

    const int FILE_HEADER_SIZE = 8+8+8+8+32;

    char header[FILE_HEADER_SIZE];
    if (fread(header, 1, FILE_HEADER_SIZE, fp) != FILE_HEADER_SIZE)
    {
        printf("Failed to read magic 8 bytes header from file\n");
        fclose(fp);
        return NULL;
    }

    if (memcmp(header, "CASCADE1", 8) != 0)
    {
        printf("File does not contain magic 8 bytes in header\n");
        fclose(fp);
        return NULL;
    }

    csc_file_t* casc_file_data = (csc_file_t*)malloc(sizeof(csc_file_t));

    casc_file_data->targetsize = be64toh(*((unsigned long long*)&header[16]));
    casc_file_data->blocksize = be64toh(*((unsigned long long*)&header[24]));
    uint8_t x[32];
    for (int i = 0; i < 32; i++) {
        x[i] = (uint8_t)(*((unsigned long long*)&header[32+i]));
    }

    csc_hashdata_t* hash = {x};
    casc_file_data->targethash = *hash;
    casc_file_data->blockcount = 1 + floor(
        (casc_file_data->targetsize - 1.0)/casc_file_data->blocksize
    );
    casc_file_data->trailblocksize = (casc_file_data->blockcount * casc_file_data->blocksize) - casc_file_data->targetsize;

    csc_block_t* block_list = malloc(
        sizeof(csc_block_t) * casc_file_data->blockcount
    );

    casc_file_data->blocks = block_list;

    for (unsigned long long b = 0;b < casc_file_data->blockcount; b++) {
        csc_block_t* block = &(block_list[b]);
        block->index = b;
        block->offset = b * casc_file_data->blocksize;
        if (b == casc_file_data->blockcount - 1 ) {
            block->length = casc_file_data->trailblocksize;
        } else {
            block->length = casc_file_data->blocksize;
        }

        block->completed = 0;

        uint8_t block_x[32];
        if (fread(block_x, 1, 32, fp) != 32) {
            printf("Cascade file not readable\n");
            fclose(fp);
            return NULL;
        }
        csc_hashdata_t* hash = {block_x};
        block->hash = *hash;
    }

    fclose(fp);

    fp = fopen(destination, "a+w");
    if (fp == NULL)
    {
        printf("Failed to open destination file %s\n", destination);
        csc_free_file(casc_file_data);
        return NULL;
    }

    void* buffer = malloc(casc_file_data->blocksize);
    if (buffer == NULL)
    {
        printf("No block buffer asigned: %d\n", casc_file_data->blocksize);
        csc_free_file(casc_file_data);
        fclose(fp);
        return NULL;
    }

    SHA256_CTX shactx;
    for(unsigned long long i = 0; i < casc_file_data->blockcount; i++)
    {
        uint8_t shabuffer[SHA256_HASH_SIZE];
        unsigned long long size = casc_file_data->blocks[i].length;
        if (fread(buffer, size, 1, fp) != 1)
        {
            break;
        }

        sha256_init(&shactx);
        sha256_update(&shactx, buffer, size);
        sha256_final(&shactx, &shabuffer);

        if (memcmp((&(&casc_file_data->blocks[i])->hash)->x, shabuffer, 32) == 0) {
            (&casc_file_data->blocks[i])->completed = 1;
        } else {
            (&casc_file_data->blocks[i])->completed = 0;
        }
    }
    fclose(fp);

    return casc_file_data;
}

/*
 * Releases the memory allocated by a file datastructure
 */
void csc_free_file(csc_file_t* file)
{
    free(file->blocks);
    file->blocks = NULL;
    free(file);
}

size_t readbytes(int sock, void* buffer, size_t count)
{
    size_t remaining = count;
    while(remaining > 0)
    {
        size_t read = recv(sock, buffer, remaining, NULL);
        if (read == 0)
            return 0;
        buffer += read;
        remaining -= read;
    }

    return count;
}

void get_block(csc_block_t* block, csc_peer_t peer, unsigned char* hash, char* output_file)
{
    printf("Attempting to get block %d from %s:%s for %s\n", block->index, peer.ip, peer.port, output_file);

    rio_t rio;
    char rio_buf[MAX_LINE];
    int peer_socket;

    /*
    TODO Request a block from a peer
    */

    FILE* fp = fopen(output_file, "rb+");
    if (fp == 0)
    {
        printf("Failed to open destination: %s\n", output_file);
        Close(peer_socket);
        return;
    }

    /*
    TODO Write the block into the data file
    */
    int write_count;

    printf("Got block %d. Wrote from %d to %d\n", block->index, block->offset, block->offset+write_count-1);
    Close(peer_socket);
    fclose(fp);
}

int get_peers_list(csc_peer_t** peers, uint8_t hash[32])
{
    rio_t rio;
    uint8_t rio_buf[MAX_LINE];
    for (int i = 0; i < MAX_LINE ; i++) {
        rio_buf[i] = 0;
    }

    int tracker_socket;

    /*
    TODO Setup a connection to the tracker
    */

    struct RequestHeader request_header;
    strncpy(request_header.protocol, "CASC", 4);
    request_header.version = htonl(1);
    request_header.command = htonl(1);
    request_header.length = htonl(BODY_SIZE);
    memcpy(rio_buf, &request_header, HEADER_SIZE);

    struct RequestBody request_body;
    strncpy(request_body.hash, hash, 32);
    inet_pton(AF_INET, my_ip, &request_body.ip);
    request_body.port = atol(my_port);
    memcpy(&rio_buf[HEADER_SIZE], &request_body, BODY_SIZE);
    for (int i = 0;i<128;i++) {
        printf("%02x ",rio_buf[i]);
    }
    printf("\n");

    return 0; // TODO REMOVE

    Rio_writen(tracker_socket, rio_buf, MESSAGE_SIZE);

    Rio_readnb(&rio, rio_buf, MAXLINE);

    char reply_header[REPLY_HEADER_SIZE];
    memcpy(reply_header, rio_buf, REPLY_HEADER_SIZE);

    uint32_t msglen = ntohl(*(uint32_t*)&reply_header[1]);
    if (msglen == 0)
    {
        Close(tracker_socket);
        return 0;
    }

    if (reply_header[0] != 0)
    {
        char* error_buf = malloc(msglen + 1);
        if (error_buf == NULL)
        {
            printf("Tracker error %d and out-of-memory reading error\n", reply_header[0]);
            Close(tracker_socket);
            return NULL;
        }
        memset(error_buf, 0, msglen + 1);
        memcpy(error_buf, &rio_buf[REPLY_HEADER_SIZE], msglen); // Fixed by Rune
        printf("Tracker gave error: %d - %s\n", reply_header[0], error_buf);
        free(error_buf);
        Close(tracker_socket);
        return NULL;
    }

    if (msglen % 12 != 0)
    {
        printf("LIST response from tracker was length %llu but should be evenly divisible by 12\n", msglen);
        Close(tracker_socket);
        return NULL;
    }

    /*
    TODO Parse the body of the response to get a list of peers

    HINT Some of the later provided code expects the peers to be stored in the ''peers' variable, which
    is an array of 'csc_peer's, as defined in cascade.h
    */
    int peercount;

    Close(tracker_socket);
    return peercount;
}

// int main() {
//     csc_file_t* f = csc_parse_file("./tests/shakespeare.10kib.txt.cascade","shakespeare.txt");
//     printf("File size: %i\nBlock Size: %i\nBlock Count: %i\nFile hash: ", f->targetsize, f->blocksize, f->blockcount);
//     for (int i = 0;i<32;i++) {
//         printf("%02x ",(&f->targethash)->x[i]);
//     }
//     printf("\nFirst block hash: ");
//     for (int i = 0;i<32;i++) {
//         printf("%02x ",(&(&f->blocks[0])->hash)->x[i]);
//     }
//     printf("\n");
// }

int main(int argc, char **argv)
{
    if (argc != MAIN_ARGNUM + 1)
    {
        fprintf(stderr, "Usage: %s <cascade file(s)> <tracker server ip> <tracker server port> <peer ip> <peer port>.\n", argv[0]);
        exit(EXIT_FAILURE);
    }
    else if (!is_valid_ip(argv[2]))
    {
        fprintf(stderr, ">> Invalid tracker IP: %s\n", argv[2]);
        exit(EXIT_FAILURE);
    }
    else if (!is_valid_port(argv[3]))
    {
        fprintf(stderr, ">> Invalid tracker port: %s\n", argv[3]);
        exit(EXIT_FAILURE);
    }
    else if (!is_valid_ip(argv[4]))
    {
        fprintf(stderr, ">> Invalid peer IP: %s\n", argv[4]);
        exit(EXIT_FAILURE);
    }
    else if (!is_valid_port(argv[5]))
    {
        fprintf(stderr, ">> Invalid peer port: %s\n", argv[5]);
        exit(EXIT_FAILURE);
    }

    snprintf(tracker_ip,   IP_LEN,   argv[2]);
    snprintf(tracker_port, PORT_LEN, argv[3]);
    snprintf(my_ip,   IP_LEN,   argv[4]);
    snprintf(my_port, PORT_LEN, argv[5]);

    char cas_str[strlen(argv[1])];
    snprintf(cas_str, strlen(argv[1])+1, argv[1]);
    char delim[] = ":";

    int casc_count = count_occurences(argv[1], ':') + 1;
    char *cascade_files[casc_count];

    char *ptr = strtok(cas_str, delim);
    int i = 0;

    while (ptr != NULL)
    {
        if (strstr(ptr, ".cascade") != NULL)
        {
            cascade_files[i++] = ptr;
            ptr = strtok(NULL, delim);
        }
        else
        {
            printf("Abort on %s\n", ptr);
            fprintf(stderr, ">> Invalid cascade file: %s\n", ptr);
            exit(EXIT_FAILURE);
        }
    }

    for (int j=0; j<casc_count; j++)
    {
        download_only_peer(cascade_files[j]);
    }

    exit(EXIT_SUCCESS);
}