/*
 * ==================================================================================
 * Server Engine - Server-side file reception and session management
 *
 * ISA Project - ICMP Covert Channel File Transfer System
 * Author: Roman Necas (xnecasr00)
 * Date: 13. October 2025
 *
 * DESCRIPTION:
 * This module implements the complete server-side logic for receiving files
 * via ICMP covert channel. It handles packet reception, session management,
 * decryption, and file writing in a continuous listening loop.
 *
 * KEY FUNCTIONS:
 * - server_run() - Main server loop (never returns until interrupted)
 * - handle_start_packet() - Process START packet and initialize session
 * - handle_data_packet() - Process DATA packet, decrypt, and write to file
 * - handle_end_packet() - Finalize transfer and send final ACK
 * - send_ack_reply() - Send ACK response to client
 * - cleanup_session() - Release session resources
 *
 * SERVER WORKFLOW:
 * 1. Initialize dual-stack sockets (IPv4 + IPv6), state machine, signal handlers
 * 2. Main Loop: Poll packets → Parse IFTP → Route to handler → Check timeouts
 * 3. Shutdown: Cleanup session → Close sockets → Exit
 *
 * SESSION MANAGEMENT:
 * Single active session tracks: session_id, expected_seq, crypto context,
 * file manager, client address. Sessions timeout after 60s of inactivity.
 *
 * PACKET HANDLING:
 * - START: Extract metadata, validate filename, init crypto/file, send ACK
 * - DATA: Validate sequence, decrypt, write to file, send ACK
 *   - Duplicate (seq < expected): Re-send ACK, don't reprocess
 *   - Out-of-order (seq > expected): Log warning, don't ACK
 * - END: Finalize decryption, close file, send ACK, cleanup session
 *
 * FILE SECURITY:
 * - Validates filenames: rejects "../" and absolute paths
 * - Extracts basename only (get_safe_filename)
 * - Handles collisions: file_1.txt, file_2.txt, etc.
 * - Files saved in current directory only
 *
 * ENCRYPTION:
 * - Uses client's IV and salt from START packet
 * - Derives same key using login + salt (must match client)
 * - AES-256-CBC with PBKDF2 (600,000 iterations)
 * - Streaming decryption (chunk-by-chunk, not entire file)
 *
 * STATE MACHINE STATES:
 * LISTENING → SESSION_INIT → RECEIVING_DATA → SESSION_CLOSING → COMPLETED
 * Error/timeout states: ERROR, TIMEOUT (return to LISTENING)
 *
 * DUAL-STACK SUPPORT:
 * - Creates two sockets (IPv4 + IPv6)
 * - Uses select() to poll both simultaneously
 * - Supports clients via IPv4 or IPv6
 * - ACKs sent using same address family as received packet
 *
 * SIGNAL HANDLING:
 * SIGINT/SIGTERM: Cleanup session and exit gracefully
 *
 * PERFORMANCE:
 * Non-blocking reception with 10ms sleep, streaming decryption, memory efficient
 * ==================================================================================
 */

#include "secret.h"

static server_session_t active_session = {0};
static server_state_machine_t server_sm = {0};

static int send_ack_reply(network_interface_t *net, uint16_t session_id, uint32_t seq_num,
                         const struct sockaddr_storage *client_addr, socklen_t client_addr_len) {
    uint8_t packet_buffer[2048];
    size_t packet_len;

    printf("Creating ACK packet for session %u, seq %u\n", session_id, seq_num);
    if (create_ack_packet(session_id, seq_num, packet_buffer, &packet_len) != 0) {
        print_error("Failed to create ACK packet");
        return -1;
    }

    printf("ACK packet created successfully, size: %zu bytes\n", packet_len);
    
    // Use network_send_reply which properly formats ICMP ECHO REPLY
    if (network_send_reply(net, packet_buffer, packet_len, client_addr, client_addr_len) != 0) {
        print_error("Failed to send ACK reply");
        return -1;
    }

    printf("Sent ACK for seq %u\n", seq_num);
    return 0;
}

static int handle_start_packet(network_interface_t *net, const iftp_header_t *header,
                             const uint8_t *payload, size_t payload_len,
                             const struct sockaddr_storage *client_addr, socklen_t client_addr_len) {
    if (payload_len != sizeof(iftp_metadata_t)) {
        print_error("Invalid START packet payload size");
        return -1;
    }

    // Clean up any existing session
    if (active_session.active) {
        printf("Warning: New session while another is active, terminating old session\n");
        file_manager_cleanup(&active_session.file_mgr);
        crypto_cleanup(&active_session.crypto);
        if (active_session.reassembly_buffer) {
            free(active_session.reassembly_buffer);
            active_session.reassembly_buffer = NULL;
        }
        if (active_session.received_fragments) {
            free(active_session.received_fragments);
            active_session.received_fragments = NULL;
        }
        memset(&active_session, 0, sizeof(active_session));
    }

    const iftp_metadata_t *metadata = (const iftp_metadata_t *)payload;

    // Initialize new session
    active_session.session_id = header->session_id;
    active_session.expected_seq = 1;  // First DATA packet will have seq=1
    active_session.file_size = metadata->file_size;
    memcpy(&active_session.client_addr, client_addr, client_addr_len);
    active_session.client_addr_len = client_addr_len;

    if (validate_filename(metadata->filename) != 0) {
        print_error("Invalid filename received - potential security risk");
        return -1;
    }

    if (get_safe_filename(metadata->filename, active_session.filename, sizeof(active_session.filename)) != 0) {
        print_error("Failed to create safe filename");
        return -1;
    }

    printf("New session started:\n");
    printf("  Session ID: %u\n", active_session.session_id);
    printf("  Filename: %s\n", active_session.filename);
    printf("  File size: %lu bytes\n", active_session.file_size);

    if (crypto_init(&active_session.crypto, login, metadata->salt, metadata->iv) != 0) {
        print_error("Failed to initialize decryption");
        return -1;
    }

    if (file_manager_open_write(&active_session.file_mgr, active_session.filename) != 0) {
        print_error("Failed to create output file");
        crypto_cleanup(&active_session.crypto);
        return -1;
    }

    file_manager_set_file_size(&active_session.file_mgr, active_session.file_size);

    // Initialize decryption context
    if (crypto_init_decrypt(&active_session.crypto) != 0) {
        print_error("Failed to initialize decryption context");
        file_manager_cleanup(&active_session.file_mgr);
        crypto_cleanup(&active_session.crypto);
        return -1;
    }

    active_session.active = 1;

    // Send ACK for START packet (seq=0)
    if (send_ack_reply(net, header->session_id, header->seq_num, client_addr, client_addr_len) == 0) {
        server_sm_handle_event(&server_sm, EVENT_PACKET_SENT, NULL);
        return 0;
    }

    return -1;
}

static int handle_data_packet(network_interface_t *net, const iftp_header_t *header,
                            const uint8_t *payload, size_t payload_len,
                            const struct sockaddr_storage *client_addr, socklen_t client_addr_len) {
    if (!active_session.active || header->session_id != active_session.session_id) {
        printf("Ignoring DATA packet from inactive or wrong session\n");
        return 0;
    }

    printf("Processing DATA packet seq %u (expected %u), %zu bytes\n", 
           header->seq_num, active_session.expected_seq, payload_len);

    if (header->seq_num == active_session.expected_seq) {
        // Process new data
        if (payload_len > 0) {
            uint8_t decrypted_buffer[MAX_PAYLOAD_SIZE + AES_BLOCK_SIZE];
            int decrypted_len;

            if (crypto_decrypt(&active_session.crypto, payload, payload_len,
                              decrypted_buffer, &decrypted_len) != 0) {
                print_error("Failed to decrypt data packet");
                return -1;
            }

            if (decrypted_len > 0) {
                if (file_manager_write_chunk(&active_session.file_mgr, decrypted_buffer, decrypted_len) != 0) {
                    print_error("Failed to write decrypted data to file");
                    return -1;
                }
                printf("Wrote %d decrypted bytes to file\n", decrypted_len);
            }
        }

        active_session.expected_seq++;
    } else if (header->seq_num < active_session.expected_seq) {
        // Duplicate packet - just re-send ACK
        printf("Duplicate DATA packet seq %u, re-sending ACK\n", header->seq_num);
    } else {
        // Out of order - don't process
        printf("Out-of-order DATA packet: expected %u, got %u\n", 
               active_session.expected_seq, header->seq_num);
        return 0;  // Don't ACK out-of-order packets
    }

    // Send ACK for this packet
    if (send_ack_reply(net, header->session_id, header->seq_num, client_addr, client_addr_len) == 0) {
        if (header->seq_num == active_session.expected_seq - 1) {
            // Only trigger state transition for new packets
            server_sm_handle_event(&server_sm, EVENT_PACKET_SENT, NULL);
        }
        return 0;
    }

    return -1;
}

static int handle_end_packet(network_interface_t *net, const iftp_header_t *header,
                           const struct sockaddr_storage *client_addr, socklen_t client_addr_len) {
    if (!active_session.active || header->session_id != active_session.session_id) {
        printf("Ignoring END packet from inactive or wrong session\n");
        return 0;
    }

    printf("Processing END packet for session %u, seq %u\n", header->session_id, header->seq_num);

    // Finalize decryption
    uint8_t final_buffer[AES_BLOCK_SIZE];
    int final_len;

    if (crypto_finalize_decrypt(&active_session.crypto, final_buffer, &final_len) != 0) {
        // This may fail if there's no padding to remove, which is OK for some cases
        printf("Note: Decryption finalization returned error (may be normal for certain file sizes)\n");
    } else if (final_len > 0) {
        if (file_manager_write_chunk(&active_session.file_mgr, final_buffer, final_len) != 0) {
            print_error("Failed to write final decrypted data");
            return -1;
        }
        printf("Wrote final %d decrypted bytes to file\n", final_len);
    }

    printf("File transfer completed successfully!\n");
    printf("  File: %s\n", active_session.filename);
    printf("  Size: %lu bytes processed\n", file_manager_get_bytes_processed(&active_session.file_mgr));

    // Send final ACK
    if (send_ack_reply(net, header->session_id, header->seq_num, client_addr, client_addr_len) == 0) {
        server_sm_handle_event(&server_sm, EVENT_PACKET_SENT, NULL);
    }

    // Clean up session
    file_manager_cleanup(&active_session.file_mgr);
    crypto_cleanup(&active_session.crypto);
    if (active_session.reassembly_buffer) {
        free(active_session.reassembly_buffer);
        active_session.reassembly_buffer = NULL;
    }
    if (active_session.received_fragments) {
        free(active_session.received_fragments);
        active_session.received_fragments = NULL;
    }
    memset(&active_session, 0, sizeof(active_session));

    return 0;
}

static void cleanup_session(void) {
    if (active_session.active) {
        file_manager_cleanup(&active_session.file_mgr);
        crypto_cleanup(&active_session.crypto);
        if (active_session.reassembly_buffer) {
            free(active_session.reassembly_buffer);
        }
        if (active_session.received_fragments) {
            free(active_session.received_fragments);
        }
        memset(&active_session, 0, sizeof(active_session));
    }
}

static void signal_handler(int sig) {
    (void)sig;
    printf("\nReceived interrupt signal, cleaning up...\n");
    cleanup_session();
    exit(0);
}

int server_run(const config_t *config) {
    if (!config) {
        print_error("Invalid configuration");
        return -1;
    }

    printf("Listening for incoming ICMP file transfers...\n");
    printf("Press Ctrl+C to stop\n");

    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);

    /* Initialize dual-stack network interface (IPv4 + IPv6) */
    network_interface_t net;
    if (network_init(&net, NULL) < 0) {
        print_error("Failed to initialize network interface");
        return -1;
    }

    server_sm_init(&server_sm);

    uint8_t recv_buffer[2048];
    struct sockaddr_storage client_addr = {0};  /* Initialize to zero to avoid uninitialized values */
    socklen_t client_addr_len = sizeof(client_addr);

    while (1) {
        /* Receive from dual-stack interface (select() handles both IPv4 and IPv6) */
        int recv_len = network_recv_packet_nonblocking(&net, recv_buffer, sizeof(recv_buffer),
                                                        &client_addr, &client_addr_len);

        if (recv_len > 0) {
            iftp_header_t header;
            uint8_t payload[MAX_PAYLOAD_SIZE];
            size_t payload_len;

            if (parse_iftp_packet(recv_buffer, recv_len, &header, payload, &payload_len) == 0) {
                char client_ip[INET6_ADDRSTRLEN];
                if (client_addr.ss_family == AF_INET) {
                    struct sockaddr_in *sin = (struct sockaddr_in *)&client_addr;
                    inet_ntop(AF_INET, &sin->sin_addr, client_ip, sizeof(client_ip));
                } else {
                    struct sockaddr_in6 *sin6 = (struct sockaddr_in6 *)&client_addr;
                    inet_ntop(AF_INET6, &sin6->sin6_addr, client_ip, sizeof(client_ip));
                }

                // Don't process ACK packets on server
                if (header.msg_type == IFTP_ACK) {
                    printf("Warning: Server received ACK packet - ignoring\n");
                    continue;
                }

                server_sm_handle_event(&server_sm, EVENT_PACKET_RECEIVED, &header);

                switch (header.msg_type) {
                    case IFTP_START:
                        if (server_sm.current_state == SERVER_LISTENING ||
                            server_sm.current_state == SERVER_SESSION_INIT) {
                            if (handle_start_packet(&net, &header, payload, payload_len,
                                                  &client_addr, client_addr_len) != 0) {
                                fprintf(stderr, "Error handling START packet\n");
                                server_sm_handle_event(&server_sm, EVENT_ERROR, NULL);
                            }
                        }
                        break;

                    case IFTP_DATA:
                        if (server_sm.current_state == SERVER_RECEIVING_DATA ||
                            server_sm.current_state == SERVER_DATA_ACK_SENT) {
                            if (handle_data_packet(&net, &header, payload, payload_len,
                                                 &client_addr, client_addr_len) != 0) {
                                fprintf(stderr, "Error handling DATA packet\n");
                                server_sm_handle_event(&server_sm, EVENT_ERROR, NULL);
                            }
                        }
                        break;

                    case IFTP_END:
                        if (server_sm.current_state == SERVER_RECEIVING_DATA ||
                            server_sm.current_state == SERVER_DATA_ACK_SENT ||
                            server_sm.current_state == SERVER_SESSION_CLOSING) {
                            if (handle_end_packet(&net, &header, &client_addr, client_addr_len) != 0) {
                                fprintf(stderr, "Error handling END packet\n");
                                server_sm_handle_event(&server_sm, EVENT_ERROR, NULL);
                            }
                        }
                        break;

                    default:
                        printf("Received unknown packet type: %d\n", header.msg_type);
                        break;
                }
            }
        }

        // Check for timeouts
        if (server_sm.current_state != SERVER_LISTENING && server_sm_check_timeout(&server_sm)) {
            printf("Session timeout detected\n");
            cleanup_session();
            server_sm_handle_event(&server_sm, EVENT_TIMEOUT, NULL);
        }

        // Reset to listening after completion
        if (server_sm.current_state == SERVER_COMPLETED) {
            printf("Session completed, returning to listening mode\n");
            server_sm_transition(&server_sm, SERVER_LISTENING, "Ready for next transfer");
        }

        usleep(10000);  // 10ms delay
    }

    cleanup_session();
    network_cleanup(&net);
    return 0;
}