#include "util.h"
#include "rlocs.h"
#include "packet.h"

#include <unistd.h>
#include <errno.h>

#include <sys/types.h>
#include <sys/wait.h>
#include <sys/uio.h>

#define MAX_SESSIONS 32

int process_icmpv4_rloc_update( struct rlocs *reg, struct packet *packet  )
{
    uint16_t hdr_len = packet->hdr.ip.ihl * 4;
    uint16_t inner_ip_hdr_offset = hdr_len + sizeof( struct icmphdr );
    if ( ntohs( packet->hdr.ip.tot_len ) < inner_ip_hdr_offset + sizeof( struct iphdr ) ) {
        debug( "Truncated ICMP packet is unidentifiable" );
        return 0;
    }

    struct rloc *s_rloc = NULL;
    struct rloc *d_rloc = NULL;
    
    struct icmphdr *icmp = (struct icmphdr *) ( packet->payload + ( packet->hdr.ip.ihl * 4 ) );

    if ( icmp->type != ICMP_DEST_UNREACH && icmp->code != ICMP_FRAG_NEEDED ) {
        return 0; // It may be going elsewhere
    }
    
    // Be careful with this - some of payload may be past allocated memory
    struct packet *inner_ip = (struct packet *) ((char *) packet + inner_ip_hdr_offset );
    
    //  Not much we can do with this case right now
    if ( inner_ip->hdr.ip.protocol != IPPROTO_HIDE_EID ) {
        debug( "ICMP Too Big response to an unwrapped packet. Peculiar." );
        return -1;
    }
    
    if ( !rlocs_find_two_ipv4( reg, &s_rloc, (struct in_addr *)&inner_ip->hdr.ip.saddr, &d_rloc, (struct in_addr *)&inner_ip->hdr.ip.daddr ) ) {
            return 0;
    }

    // All we're interested in is setting path mtu
    uint16_t new_mtu = ntohs( icmp->un.frag.mtu );
    rlocs_set_path_mtu( reg, s_rloc, d_rloc, new_mtu );
    debug( "Set MTU for %s <-> %s to %u", s_rloc->presentation, d_rloc->presentation, new_mtu );
    return 1;
}


int process_icmpv6_rloc_update( struct rlocs *reg, struct packet *packet )
{
    warn( "STUB: process_icmpv6_rloc_update" );
    return 0;
}

int process_packet( struct rlocs *reg, struct packet *pkt, struct rsp_data *frag1, struct rsp_data *frag2 )
{
    int protocol = packet_find_protocol( pkt );

    if ( protocol == -1 ) {
        warn( "Couldn't work out version / protocol of received packet, discarding" );
        return 0;
    }

    if ( protocol == IPPROTO_HIDE_EID ) {
        return unwrap_packet( reg, pkt, frag1 );
    }

    if ( protocol == IPPROTO_ICMP && process_icmpv4_rloc_update( reg, pkt ) ) {
        return 0;
    }

    if ( protocol == IPPROTO_ICMPV6 && process_icmpv6_rloc_update( reg, pkt ) ) {
        return 0;
    }

    return wrap_packet( reg, pkt, frag1, frag2 );
}


#define ARG_NUM_SESSIONS  1
#define ARG_RLOC_DB       2
#define ARG_IFNAME        3
#define ARG_LAST          4

int do_session( int multi, int argc, char **argv )
{
    struct packet recv_pkt;
    struct session session;
    
    struct rsp_data frag1;
    struct rsp_data frag2;
    ssize_t count;

    rlocs_init();

    if ( !session_setup( &session, argv[ARG_RLOC_DB], argv[ARG_IFNAME], multi ) ) {
        warn( "Failed to set up session, exiting" );
        return 1;
    }

    if  (argc > ARG_LAST ) {
        if ( !session_upgrade_rlocs( &session, argc - ARG_LAST, argv + ARG_LAST ) ) {
            warn( "Failed to upgrade rlocs for session, exiting" );
            session_teardown( &session );
            return 1;
        }    
    }    
    
    memset( &recv_pkt, 0, sizeof( struct packet ) );
    memset( &frag1, 0, sizeof( struct rsp_data ) );
    memset( &frag2, 0, sizeof( struct rsp_data ) );

    warn( "TODO: Write BGP interventions to file" );

    info( "Listening for packets" );
    while( 1 ) {
        int rsp_count;
        
        if ( ( count = read( session.fd, &recv_pkt, sizeof( struct packet ) ) ) < 0 ) {
            warn( "Failed to get a packet (%s)", strerror( errno ) );
            break;
        }
        
        if ( count == 0 ) {
            warn( "Got EOF" );
            break;
            
        }

        if ( ( rsp_count = process_packet( session.rlocs, &recv_pkt, &frag1, &frag2 ) ) < 0 ) {
            debug( "Error processing packet, dropping" );
        }
        
        if ( rsp_count > 0 ) {
            if ( ( count = writev( session.fd, frag1.iovs, frag1.count ) ) < 0 ) {
                debug( "Error writing processed packet to output: %s", strerror(errno) );
            }
        }
        
        if ( rsp_count == 2 ) {
            if ( ( count = writev( session.fd, frag2.iovs, frag2.count ) ) < 0 ) {
                debug( "Error writing second processed packet to output: %s", strerror(errno) );
            }
        }
        
    }
    
    info( "Finished, cleaning up" );
    session_teardown( &session );

    return 0;
}

/*
 * Entry point. Expects an invocation like:
 * hide-eid <filename of rloc database> <listen_ifname> <output_ifname>
 * 
 * hide-eid wraps packets that come to it unwrapped, and unwraps packets that
 * come to it wrapped. This makes it handy for all-in-one boxes
 */ 
int main(int argc, char **argv)
{
    pid_t pids[MAX_SESSIONS];
    int num_sessions, i;
    int num_returned = 0, ret = 0;

    if ( argc < ARG_LAST ) {
        warn( "Usage: %s <number of sessions> <rloc database> <ifname> [ <rloc> <keyfile> ]n", argv[0] );
        return 1;
    }
    
    memset( &pids, 0, sizeof( pid_t ) * MAX_SESSIONS );

    num_sessions = atoi( argv[ARG_NUM_SESSIONS] );

    if ( num_sessions <= 0 || num_sessions > MAX_SESSIONS ) {
        warn( "Bad value: %s for number of sessions", argv[ARG_NUM_SESSIONS] );
    }

    // Don't spawn a child process if there's only 1 of us
    if ( num_sessions == 1 ) {
        return do_session( 0, argc, argv );
    }

    for( i = 0 ;  i < num_sessions ; i++ ) {
        pids[i] = fork();
        if ( pids[i] == -1 ) {
            warn( "fork() failure!" );
            exit(1);
        }

        if ( pids[i] == 0 ) {
            int ret = do_session( 1, argc, argv );
            info( "Child process %i exiting with status %i", getpid(), ret );
            exit( ret );
        } else {
            info( "Child session %i started", pids[i] );
        }
        
    }

    // wait on all sessions returning
    // todo: pass sigterm, etc to children
    while ( num_returned < num_sessions ) {
        int status;
        pid_t pid;

        pid = waitpid( -1, &status, 0 );
        if ( pid == -1 ) {
            warn( "waitpid() failure: %s", strerror(errno) );
            ret = 1;
            break;
        }

        if ( WIFEXITED( status ) ) {
            for ( i = 0 ; i < num_sessions ; i++ ) {
                if ( pids[i] == pid ) {
                    info( "Child pid %i exited", pid );
                    pids[i] = -1;
                    num_returned += 1;
                    // TODO: if exit status was bad, set ret to 1
                    break;
                }
            }
        }
    }
    
    return ret;

}