/*
 * Client Engine - Client-side file transfer logic
 * Author: Roman Necas (xnecasr00)
 * Date: 13. October 2025
 *
 * This module implements the client-side file transfer operations:
 * - File encryption and fragmentation
 * - ICMP packet transmission
 * - Acknowledgment handling
 * - State machine-driven transfer flow
 */

#include "secret.h"

/*
 * Wait for ACK packet from server (non-blocking)
 * Returns 0 if matching ACK received, -1 otherwise
 */
static int wait_for_ack_nonblocking(network_interface_t *net, uint16_t expected_session_id, uint32_t expected_seq) {
    uint8_t recv_buffer[2048];
    struct sockaddr_storage src_addr;
    socklen_t src_addr_len = sizeof(src_addr);

    /* Try to receive a packet (non-blocking, returns immediately if nothing available) */
    int recv_len = network_recv_packet_nonblocking(net, recv_buffer, sizeof(recv_buffer), &src_addr, &src_addr_len);

    /* No packet received or error */
    if (recv_len <= 0) {
        return -1;
    }

    iftp_header_t header;
    uint8_t payload[MAX_PAYLOAD_SIZE];
    size_t payload_len;

    /* Parse the IFTP packet and validate it */
    if (parse_iftp_packet(recv_buffer, recv_len, &header, payload, &payload_len) == 0) {
        /* Check if this is the ACK we're waiting for */
        if (header.session_id == expected_session_id &&
            header.msg_type == IFTP_ACK &&
            header.seq_num == expected_seq) {
            return 0;  /* Success - got the ACK we wanted */
        }
    }

    return -1;  /* Not the ACK we're waiting for */
}

/*
 * Send START packet to initiate file transfer
 * START packet contains: filename, file size, IV, and salt for encryption
 */
static int send_start_packet(network_interface_t *net, crypto_engine_t *crypto,
                           file_manager_t *fm, uint16_t session_id) {
    iftp_metadata_t metadata = {0};

    /* Extract basename from full path (security: only send filename, not full path) */
    if (file_manager_get_basename(file_manager_get_filename(fm), metadata.filename,
                                sizeof(metadata.filename)) != 0) {
        print_error("Filename too long or invalid");
        return -1;
    }

    /* Fill in metadata: file size and crypto parameters */
    metadata.file_size = file_manager_get_file_size(fm);
    crypto_get_iv(crypto, metadata.iv);      /* Get IV for this transfer */
    crypto_get_salt(crypto, metadata.salt);  /* Get salt for key derivation */

    uint8_t packet_buffer[2048];
    size_t packet_len;

    /* Create IFTP START packet with metadata as payload */
    if (create_start_packet(session_id, &metadata, packet_buffer, &packet_len) != 0) {
        print_error("Failed to create START packet");
        return -1;
    }

    /* Send START packet via ICMP */
    if (network_send_packet(net, packet_buffer, packet_len) != 0) {
        print_error("Failed to send START packet");
        return -1;
    }

    printf("Sent START packet (file: %s, size: %lu bytes)\n",
           metadata.filename, metadata.file_size);

    return 0;
}

/*
 * Encrypt file data and send as DATA packet
 * Note: AES-CBC may buffer data, so encrypted length can be 0 (data buffered in crypto context)
 */
static int send_data_packet_with_encryption(network_interface_t *net, crypto_engine_t *crypto,
                                           file_manager_t *fm, uint16_t session_id, uint32_t seq_num,
                                           uint8_t *file_buffer, size_t bytes_read) {
    uint8_t encrypted_buffer[MAX_PAYLOAD_SIZE + AES_BLOCK_SIZE];
    uint8_t packet_buffer[2048];
    size_t packet_len;

    /* Encrypt the file data chunk using AES-256-CBC */
    int enc_len;
    if (crypto_encrypt(crypto, file_buffer, bytes_read, encrypted_buffer, &enc_len) != 0) {
        print_error("Failed to encrypt data chunk");
        return -1;
    }

    /* EVP_EncryptUpdate may buffer data and return 0 bytes
     * This is normal for AES-CBC block cipher - data will be in the encryption context
     * Only send packet if we have actual encrypted data to send */
    if (enc_len > 0) {
        uint32_t current_frag, total_frags;
        file_manager_get_fragment_info(fm, &current_frag, &total_frags);

        /* Create DATA packet with encrypted payload */
        if (create_data_packet(session_id, seq_num, current_frag, total_frags,
                             encrypted_buffer, enc_len, packet_buffer, &packet_len) != 0) {
            print_error("Failed to create DATA packet");
            return -1;
        }

        /* Send DATA packet via ICMP */
        if (network_send_packet(net, packet_buffer, packet_len) != 0) {
            print_error("Failed to send DATA packet");
            return -1;
        }

        printf("Sent DATA packet %u/%u (seq: %u, %d bytes encrypted)\n",
               current_frag, total_frags, seq_num, enc_len);
        return enc_len;  /* Return encrypted length for statistics */
    }

    return 0;  /* No data sent yet (buffered in crypto context) */
}

/*
 * Send END packet to signal transfer completion
 */
static int send_end_packet_simple(network_interface_t *net, uint16_t session_id, uint32_t seq_num) {
    uint8_t packet_buffer[2048];
    size_t packet_len;

    /* Create END packet (no payload) */
    if (create_end_packet(session_id, seq_num, packet_buffer, &packet_len) != 0) {
        print_error("Failed to create END packet");
        return -1;
    }

    /* Send END packet via ICMP */
    if (network_send_packet(net, packet_buffer, packet_len) != 0) {
        print_error("Failed to send END packet");
        return -1;
    }

    printf("Sent END packet (seq: %u)\n", seq_num);
    return 0;
}

/*
 * Main client function - orchestrates complete file transfer process
 * State machine driven: IDLE → START_SENT → DATA_TRANSFER → DATA_WAIT_ACK → END_SENT → COMPLETED
 */
int client_run(const config_t *config) {
    if (!config) {
        print_error("Invalid configuration");
        return -1;
    }

    printf("Source file: %s\n", config->source_file);
    printf("Destination: %s\n", config->dest_address);

    /* Initialize all components */
    network_interface_t net;
    crypto_engine_t crypto;
    file_manager_t fm;
    client_state_machine_t sm;
    int result = -1;
    int data_sent = 0;
    int encryption_finalized = 0;

    memset(&net, 0, sizeof(net));
    memset(&crypto, 0, sizeof(crypto));
    memset(&fm, 0, sizeof(fm));

    /* Setup network interface with destination address */
    if (network_init(&net, config->dest_address) != 0) {
        print_error("Failed to initialize network interface");
        goto cleanup;
    }

    /* Open source file for reading */
    if (file_manager_open_read(&fm, config->source_file) != 0) {
        print_error("Failed to open source file");
        goto cleanup;
    }

    /* Initialize crypto engine - generates random IV and salt */
    if (crypto_init(&crypto, login, NULL, NULL) != 0) {
        print_error("Failed to initialize cryptographic engine");
        goto cleanup;
    }

    /* Generate cryptographically secure session ID */
    uint16_t session_id = generate_secure_session_id();
    client_sm_init(&sm, session_id);

    printf("Session ID: %u\n", session_id);

    /* Initialize encryption context once before the loop (must be done before any encryption) */
    if (crypto_init_encrypt(&crypto) != 0) {
        print_error("Failed to initialize encryption context");
        goto cleanup;
    }

    /* Main transfer loop - runs until COMPLETED, ERROR, or ABORT state reached */
    while (sm.current_state != CLIENT_COMPLETED &&
           sm.current_state != CLIENT_ERROR &&
           sm.current_state != CLIENT_ABORT) {

        /* State machine: handle current state */
        switch (sm.current_state) {
            case CLIENT_IDLE:  /* Initial state - send START packet */
                if (send_start_packet(&net, &crypto, &fm, session_id) == 0) {
                    client_sm_handle_event(&sm, EVENT_PACKET_SENT, NULL);
                } else {
                    client_sm_handle_event(&sm, EVENT_ERROR, NULL);
                }
                break;

            case CLIENT_START_SENT:  /* Waiting for START-ACK from server */
                if (client_sm_check_timeout(&sm)) {
                    /* Timeout - trigger retry with exponential backoff */
                    client_sm_handle_event(&sm, EVENT_TIMEOUT, NULL);
                    if (sm.current_state == CLIENT_START_SENT) {
                        send_start_packet(&net, &crypto, &fm, session_id);  /* Retransmit START */
                    }
                } else if (wait_for_ack_nonblocking(&net, session_id, 0) == 0) {
                    /* Got START-ACK! Proceed to data transfer */
                    client_sm_handle_event(&sm, EVENT_ACK_RECEIVED, NULL);
                    sm.current_seq = 1;  /* DATA packets start at sequence 1 (0 is for START) */
                }
                break;

            case CLIENT_DATA_TRANSFER:  /* Reading file and sending DATA packets */
                {
                    /* Read chunk - leave room for AES padding (up to 16 bytes for AES-128 blocks) */
                    const size_t chunk_size = MAX_PAYLOAD_SIZE - AES_BLOCK_SIZE;  /* 1400 - 16 = 1384 bytes */
                    uint8_t file_buffer[MAX_PAYLOAD_SIZE];
                    size_t bytes_read;

                    int read_result = file_manager_read_chunk(&fm, file_buffer, chunk_size, &bytes_read);

                    if (read_result < 0) {
                        print_error("Failed to read file chunk");
                        client_sm_handle_event(&sm, EVENT_ERROR, NULL);
                        break;
                    }

                    if (read_result == 0 || bytes_read == 0) {
                        /* End of file reached - finalize encryption and send any remaining data */
                        if (!encryption_finalized) {
                            /* Get final encrypted block with PKCS#7 padding */
                            uint8_t final_buffer[MAX_PAYLOAD_SIZE + AES_BLOCK_SIZE];
                            int final_len;

                            if (crypto_finalize_encrypt(&crypto, final_buffer, &final_len) != 0) {
                                print_error("Failed to finalize encryption");
                                client_sm_handle_event(&sm, EVENT_ERROR, NULL);
                                break;
                            }

                            encryption_finalized = 1;

                            if (final_len > 0) {
                                // Send final encrypted block (padding)
                                uint8_t packet_buffer[2048];
                                size_t packet_len;
                                uint32_t current_frag, total_frags;
                                file_manager_get_fragment_info(&fm, &current_frag, &total_frags);

                                if (create_data_packet(session_id, sm.current_seq, current_frag, total_frags,
                                                     final_buffer, final_len, packet_buffer, &packet_len) == 0 &&
                                    network_send_packet(&net, packet_buffer, packet_len) == 0) {
                                    printf("Sent final DATA packet (seq: %u, %d bytes)\n", sm.current_seq, final_len);
                                    client_sm_handle_event(&sm, EVENT_PACKET_SENT, NULL);
                                    data_sent = 1;
                                } else {
                                    print_error("Failed to send final encrypted block");
                                    client_sm_handle_event(&sm, EVENT_ERROR, NULL);
                                }
                            } else {
                                // No padding needed, trigger EOF to allow END packet
                                client_sm_handle_event(&sm, EVENT_FILE_EOF, NULL);
                            }
                        } else {
                            // Already finalized, just trigger EOF
                            client_sm_handle_event(&sm, EVENT_FILE_EOF, NULL);
                        }
                    } else {
                        // We have file data to encrypt and send
                        int enc_result = send_data_packet_with_encryption(&net, &crypto, &fm,
                                                                         session_id, sm.current_seq,
                                                                         file_buffer, bytes_read);
                        if (enc_result > 0) {
                            client_sm_handle_event(&sm, EVENT_PACKET_SENT, NULL);
                            data_sent = 1;
                        } else if (enc_result == 0) {
                            // Data was buffered in crypto context, continue reading
                            continue;
                        } else {
                            client_sm_handle_event(&sm, EVENT_ERROR, NULL);
                        }
                    }
                }
                break;

            case CLIENT_DATA_WAIT_ACK:
                if (client_sm_check_timeout(&sm)) {
                    client_sm_handle_event(&sm, EVENT_TIMEOUT, NULL);
                } else if (wait_for_ack_nonblocking(&net, session_id, sm.current_seq) == 0) {
                    printf("Received ACK for packet %u\n", sm.current_seq);
                    sm.current_seq++;  // Increment sequence after successful ACK
                    client_sm_handle_event(&sm, EVENT_ACK_RECEIVED, NULL);
                }
                break;

            case CLIENT_END_SENT:
                if (client_sm_check_timeout(&sm)) {
                    client_sm_handle_event(&sm, EVENT_TIMEOUT, NULL);
                    if (sm.current_state == CLIENT_END_SENT) {
                        send_end_packet_simple(&net, session_id, sm.current_seq);
                    }
                } else if (wait_for_ack_nonblocking(&net, session_id, sm.current_seq) == 0) {
                    printf("Received final ACK - transfer complete!\n");
                    client_sm_handle_event(&sm, EVENT_ACK_RECEIVED, NULL);
                }
                break;

            case CLIENT_COMPLETED:
            case CLIENT_ERROR:
            case CLIENT_ABORT:
                break;
        }

        // Handle EOF transition - only send END after encryption is finalized
        if (sm.current_state == CLIENT_DATA_TRANSFER && file_manager_is_complete(&fm) &&
            encryption_finalized && data_sent) {
            // END packet uses current sequence (which is already next after last ACK'd DATA)
            if (send_end_packet_simple(&net, session_id, sm.current_seq) == 0) {
                client_sm_transition(&sm, CLIENT_END_SENT, "File transfer complete");
            } else {
                client_sm_handle_event(&sm, EVENT_ERROR, NULL);
            }
        }

        usleep(10000);
    }

    if (sm.current_state == CLIENT_COMPLETED) {
        printf("File transfer completed successfully!\n");
        result = 0;
    } else {
        printf("Transfer failed in state: %s\n", client_state_name(sm.current_state));
        result = -1;
    }

cleanup:
    network_cleanup(&net);
    crypto_cleanup(&crypto);
    file_manager_cleanup(&fm);

    return result;
}