#include "mozen_tunnel.h"
#include <string.h>
// ·Ö°ü

/*
* Initialize tunnel
* tunnel: tunnel context
* prot: protocol context
* mtu: maximum transmission unit
* user_context: user context
* returns: none
*/
void mozen_tunnel_init(mozen_tunnel_t *tunnel, mozen_protocol_t *prot, uint16_t mtu, void *user_context) 
{
    if (!tunnel) return;
    memset(tunnel, 0, sizeof(mozen_tunnel_t));
    tunnel->prot = prot;
    tunnel->mtu = mtu;
    tunnel->user_context = user_context;
}

/*
 * Send a response to the tunnel command
 * tunnel: tunnel context
 * target_cmd: command id
 * is_ok: true if the command was successful, false otherwise
 * returns: none
*/
void mozen_tunnel_set_callbacks(mozen_tunnel_t *tunnel, 
                                mozen_tunnel_on_stream_start_t on_start,
                                mozen_tunnel_on_stream_data_t on_data,
                                mozen_tunnel_on_stream_finish_t on_finish) 
{
    if (!tunnel) return;
    tunnel->on_start = on_start;
    tunnel->on_data = on_data;
    tunnel->on_finish = on_finish;
}

/*
* Send a response to the tunnel command
* tunnel: tunnel context
* target_cmd: command id
* is_ok: true if the command was successful, false otherwise
* returns: none
*/
static void mozen_tunnel_send_response(mozen_tunnel_t *tunnel, uint8_t target_cmd, bool is_ok) 
{
    if (!tunnel || !tunnel->prot) return;
    
    uint8_t res_buf[sizeof(mozen_tunnel_header_t) + 2];
    mozen_tunnel_header_t *res_header = (mozen_tunnel_header_t *)res_buf;
    
    res_header->tunnel_type = TUNNEL_TYPE_DATA_TRANSFER;
    res_header->target_cmd = target_cmd;
    res_buf[sizeof(mozen_tunnel_header_t)] = is_ok ? 'o' : 'N';
    res_buf[sizeof(mozen_tunnel_header_t) + 1] = is_ok ? 'k' : 'G';
    // tunnel messages go through IRQ buffer context to avoid disrupting main buffer?
    // The previous implementation used the default send context.
    mozen_protocol_send_frame(tunnel->prot, MOZEN_TUNNEL_COMMAND_ID, res_buf, sizeof(res_buf), true);
}

/*
* Handle a received tunnel frame
* tunnel: tunnel context
* data: received data
* length: received data length
* returns: none
*/
static void mozen_tunnel_abort(mozen_tunnel_t *tunnel) 
{
    if (tunnel->is_active) 
    {
        if (tunnel->on_finish) 
        {
            tunnel->on_finish(tunnel->user_context, false);
        }
        tunnel->is_active = false;
    }
}

/*
* Handle a received tunnel frame
* tunnel: tunnel context
* data: received data
* length: received data length
* returns: none
*/
void mozen_tunnel_handle_command(mozen_tunnel_t *tunnel, const uint8_t *data, uint16_t length) 
{
    if (!tunnel || length < sizeof(mozen_tunnel_header_t)) return;

    const mozen_tunnel_header_t *tunnel_header = (const mozen_tunnel_header_t *)data;
    const uint8_t *payload = data + sizeof(mozen_tunnel_header_t);
    uint16_t payload_len = length - sizeof(mozen_tunnel_header_t);

    switch (tunnel_header->tunnel_type) 
    {
        case TUNNEL_TYPE_HANDSHAKE:  // handshake request
        {
            if (payload_len < sizeof(uint32_t) + sizeof(uint16_t))  return; 
            
            mozen_tunnel_abort(tunnel); // Abort any ongoing stream

            tunnel->is_active = true;
            tunnel->target_cmd = tunnel_header->target_cmd;
            
            memcpy(&tunnel->total_data_len, payload, sizeof(uint32_t));
            memcpy(&tunnel->total_packets, payload + sizeof(uint32_t), sizeof(uint16_t));
            
            tunnel->next_packet_index = 1;
            tunnel->received_data_len = 0;

            // Prepare and send handshake response
            uint8_t res_buf[sizeof(mozen_tunnel_handshake_res_t)];
            mozen_tunnel_handshake_res_t *res = (mozen_tunnel_handshake_res_t *)res_buf;
            res->header.tunnel_type = TUNNEL_TYPE_RESPONSE;
            res->header.target_cmd = tunnel_header->target_cmd;
            res->status = 0x6B6F; // 'ok' (little-endian: 'o' is 0x6F, 'k' is 0x6B)
            res->mtu = tunnel->mtu; 
            
            if (tunnel->prot) 
            {
                mozen_protocol_send_frame(tunnel->prot, MOZEN_TUNNEL_COMMAND_ID, (uint8_t *)res, sizeof(mozen_tunnel_handshake_res_t), true);
            }
            break;
        }

        case TUNNEL_TYPE_DATA_TRANSFER: // data transfer
        {
            if (!tunnel->is_active || tunnel_header->target_cmd != tunnel->target_cmd) return;
            if (payload_len < sizeof(uint16_t)) return;

            uint16_t packet_index;
            memcpy(&packet_index, payload, sizeof(uint16_t));

            if (packet_index != tunnel->next_packet_index) 
            {
                mozen_tunnel_send_response(tunnel, tunnel->target_cmd, false);
                mozen_tunnel_abort(tunnel);
                return;
            }

            const uint8_t *chunk_data = payload + sizeof(uint16_t);
            uint16_t chunk_len = payload_len - sizeof(uint16_t);

            bool accepted = true;

            if (packet_index == 1) 
            {
                if (tunnel->on_start) 
                {
                    accepted = tunnel->on_start(tunnel->user_context, tunnel->target_cmd, chunk_data, chunk_len);
                }
            } 
            else 
            {
                if (tunnel->on_data) 
                {
                    accepted = tunnel->on_data(tunnel->user_context, chunk_data, chunk_len);
                }
            }

            if (!accepted) 
            {
                mozen_tunnel_send_response(tunnel, tunnel->target_cmd, false);
                mozen_tunnel_abort(tunnel);
                return;
            }

            tunnel->received_data_len += chunk_len;
            tunnel->next_packet_index++;

            mozen_tunnel_send_response(tunnel, tunnel->target_cmd, true);

            if (tunnel->next_packet_index > tunnel->total_packets) 
            {
                if (tunnel->received_data_len == tunnel->total_data_len) 
                {
                    if (tunnel->on_finish) 
                    {
                        tunnel->on_finish(tunnel->user_context, true);
                    }
                } 
                else 
                {
                    if (tunnel->on_finish) 
                    {
                        tunnel->on_finish(tunnel->user_context, false);
                    }
                }
                tunnel->is_active = false;
            }
            break;
        }
    }
}