/* statquery.cc
****************************************************************************************************************
****************************************************************************************************************

    Copyright (C) 2004, 2005, 2006, 2007, 2008, 2009, 2010, 2011, 2012, 2013, 2014, 2015, 2020, 2021 Askar
    Almatov

    This file is part of Trafficmeter.

    Trafficmeter is free software: you can redistribute it and/or modify it under the terms of the GNU General
    Public License as published by the Free Software Foundation, either version 3 of the License, or (at your
    option) any later version.

    Trafficmeter is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the
    implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
    License for more details.

    You should have received a copy of the GNU General Public License along with Trafficmeter. If not, see
    <http://www.gnu.org/licenses/>.
*/

#include <cassert>
#include <iostream>
#include <fcntl.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/mman.h>

#include "fileheader.h"
#include "statquery.h"
#include "statrecord.h"
#include "tmdefs.h"

using namespace iclass;
using namespace trafficmeter;

using std::map;
using std::max;
using std::min;
using std::runtime_error;
using std::string;

struct StatQuery::Search
{
    StatQuery*                              query;
    std::map<StatQuery::Key, SummaryIo>*    rows;

    virtual         ~Search() {}
    void            extract( const StatRecord* record, time_t group );
    virtual void    find( const StatRecord* begin, const StatRecord* end, time_t group, int hash ) = 0;
};

struct SearchByIp_ : public StatQuery::Search
{
                    SearchByIp_( StatQuery* query, std::map<StatQuery::Key, SummaryIo>* rows );
    virtual void    find( const StatRecord* begin, const StatRecord* end, time_t group, int hash ) override;
};

struct SearchByNetwork_ : public StatQuery::Search
{
                    SearchByNetwork_( StatQuery* query, std::map<StatQuery::Key, SummaryIo>* rows );
    virtual void    find( const StatRecord* begin, const StatRecord* end, time_t group, int hash ) override;
};

struct SearchBySet_ : public StatQuery::Search
{
                    SearchBySet_( StatQuery* query, std::map<StatQuery::Key, SummaryIo>* rows );
    virtual void    find( const StatRecord* begin, const StatRecord* end, time_t group, int hash ) override;
};

/**************************************************************************************************************/
inline static int
compareIp_( const void* lhs, const void* rhs )
{
    return memcmp( lhs, rhs, 4 );
}

/**************************************************************************************************************/
inline static int
compareNetwork_( const void* lhs, const void* rhs )
{
    uint32_t    mask = *( reinterpret_cast<const uint32_t*>(lhs) + 1 );
    uint32_t    rhsNetwork = *reinterpret_cast<const uint32_t*>( rhs ) & mask;

    return memcmp( lhs, &rhsNetwork, 4 );
}

/**************************************************************************************************************/
StatQuery::StatQuery() :
    TimeQuery( STAT_DIR, PREFIX_HASH_LENGTH ),
    groupByTime_( nullptr ),
    shouldSaveVlans_( false ),
    shouldSaveHosts_( false )
{
}

/**************************************************************************************************************/
StatQuery::~StatQuery()
{
    if ( groupByTime_ )
    {
        delete groupByTime_;
    }
}

/**************************************************************************************************************/
bool
StatQuery::timeByHour()
{
    if ( groupByTime_ )
    {
        return false;
    }

    groupByTime_ = new GroupByHour;
    return true;
}

/**************************************************************************************************************/
bool
StatQuery::timeByDay()
{
    if ( groupByTime_ )
    {
        return false;
    }

    groupByTime_ = new GroupByDay;
    return true;
}

/**************************************************************************************************************/
bool
StatQuery::timeByMonth()
{
    if ( groupByTime_ )
    {
        return false;
    }

    groupByTime_ = new GroupByMonth;
    return true;
}

/**************************************************************************************************************/
bool
StatQuery::hideTime()
{
    if ( groupByTime_ )
    {
        return false;
    }

    groupByTime_ = new IgnoreTime;
    return true;
}

/**************************************************************************************************************/
void
StatQuery::saveVlans()
{
    shouldSaveVlans_ = true;
}

/**************************************************************************************************************/
void
StatQuery::saveHosts()
{
    shouldSaveHosts_ = true;
}

/**************************************************************************************************************/
void
StatQuery::addNetwork( ipnet4t network )
{
    networks_.insert( network );

    if ( network.networkbits() < PREFIX_HASH_LENGTH )
    {
        int     hash = network.front() >> ( 32 - PREFIX_HASH_LENGTH );
        int     nHashes = 1U << ( PREFIX_HASH_LENGTH - network.networkbits() );

        while ( nHashes-- )
        {
            hashes_.insert( hash++ );
        }
    }
    else
    {
        hashes_.insert( network.front() >> (32 - PREFIX_HASH_LENGTH) );
    }
}

/**************************************************************************************************************/
void
StatQuery::fetch( map<Key,SummaryIo>& summaries )
{
    time_t      qStart = begin().unixtime();
    time_t      qEnd = end().unixtime();
    Search*     search;

    if ( qStart >= qEnd || networks_.empty() )
    {
        return;
    }

    if ( networks_.size() > 1 )
    {
        search = new SearchBySet_( this, &summaries );
    }
    else if ( networks_.begin()->nodebits() )
    {
        search = new SearchByNetwork_( this, &summaries );
    }
    else
    {
        search = new SearchByIp_( this, &summaries );
    }

    readDirectory( hashes_ );

    TimeQuery::Range    files = getFiles( qEnd );

    for ( auto file = files.first; file != files.second; ++file )
    {
        if ( file->isScanned )
        {
            continue;
        }

        FileHeader  header;
        time_t      fStart;
        time_t      fEnd;
        long        mapOffset;

        if ( file->fd < 0 )
        {
            Path&&  path( filePath(*file) );

            file->fd = open( path.c_str(), O_RDONLY );

            if ( file->fd == -1 || !header.read(file->fd) )
            {
                throw runtime_error( "file read error: " + path + " - " + strerror(errno) );
            }

            if ( !header.isCompatible(FileHeader::Type::STAT) )
            {
                throw runtime_error( "incompatible file format: " + path );
            }

            fStart = max( moduloRoundDown(qStart, GROUP_FACTOR), header.oldestGroup() );
            fEnd = min( moduloRoundDown(qEnd-1, GROUP_FACTOR), header.newestGroup() ) + GROUP_FACTOR;
            mapOffset = header.groupOffset( fStart );
            file->mapSize = header.groupOffset( fEnd ) - mapOffset;

            if ( file->mapSize <= 0 )
            {
                file->isScanned = true;
                continue;
            }

            file->offset = mapOffset % getpagesize();
            mapOffset -= file->offset;
            file->mapSize += file->offset;
            file->mapping = static_cast<uint8_t*>( mmap(nullptr, file->mapSize, PROT_READ, MAP_SHARED, file->fd, mapOffset) );

            if ( file->mapping == MAP_FAILED )
            {
                throw runtime_error( "mmap error: " + path + " - " + strerror(errno) );
            }
        }

        const StatRecord*   records = reinterpret_cast<const StatRecord*>( file->mapping + file->offset );

        for ( time_t g = fStart; g < fEnd; g += GROUP_FACTOR )
        {
            long                nextOffset = header.groupOffset( g + GROUP_FACTOR ) - mapOffset;
            const StatRecord*   nextRecords = reinterpret_cast<const StatRecord*>( file->mapping + nextOffset );

            search->find( records, nextRecords, regroup(g), file->hash() );
            records = nextRecords;
        }

        file->isScanned = true;
    }

    for ( auto file = files.first; file != files.second; ++file )
    {
        if ( file->isScanned )
        {
            if ( file->mapping != nullptr )
            {
                munmap( file->mapping, file->mapSize );
            }

            if ( file->fd >= 0 )
            {
                close( file->fd );
            }
        }
    }

    delete search;
}

/**************************************************************************************************************/
void
StatQuery::Search::extract( const StatRecord* record, time_t group )
{
    int     vlan = ntohs( record->beStat.vlan );

    if ( !query->checkVlan(vlan) )
    {
        return;
    }

    StatQuery::Key  key;

    key.group = group;
    key.vlan = query->shouldSaveVlans()? vlan : 0;

    if ( query->shouldSaveHosts() )
    {
        key.ip = ntohl( record->beStat.ip );
    }

    (*rows)[ key ] += record->summaryIo;
}

/**************************************************************************************************************/
SearchByIp_::SearchByIp_( StatQuery* query, std::map<StatQuery::Key, SummaryIo>* rows )
{
    Search::query = query;
    Search::rows = rows;
}

/**************************************************************************************************************/
void
SearchByIp_::find( const StatRecord* begin, const StatRecord* end, time_t group, int hash )
{
    const StatRecord*   forward;
    const StatRecord*   reverse;
    uint32_t            key;

    key = htonl( query->networks().begin()->front() );
    reverse = static_cast<const StatRecord*>( bsearch(&key, begin, end-begin, sizeof(*end), compareIp_) );

    if ( reverse )
    {
        forward = reverse--;

        while ( reverse >= begin && reverse->beStat.ip == key )
        {
            extract( reverse--, group );
        }

        while ( forward < end && forward->beStat.ip == key )
        {
            extract( forward++, group );
        }
    }
}

/**************************************************************************************************************/
SearchByNetwork_::SearchByNetwork_( StatQuery* query, std::map<StatQuery::Key, SummaryIo>* rows )
{
    Search::query = query;
    Search::rows = rows;
}

/**************************************************************************************************************/
void
SearchByNetwork_::find( const StatRecord* begin, const StatRecord* end, time_t group, int hash )
{
    const StatRecord*   forward;
    const StatRecord*   reverse;
    uint32_t            key[ 2 ];

    key[ 0 ] = htonl( query->networks().begin()->front() );
    key[ 1 ] = htonl( query->networks().begin()->mask() );
    reverse = static_cast<const StatRecord*>( bsearch(key, begin, end-begin, sizeof(*end), compareNetwork_) );

    if ( reverse )
    {
        forward = reverse--;

        while ( reverse >= begin && !((reverse->beStat.ip&key[1]) ^ key[0]) )
        {
            extract( reverse--, group );
        }

        while ( forward < end && !((forward->beStat.ip&key[1]) ^ key[0]) )
        {
            extract( forward++, group );
        }
    }
}

/**************************************************************************************************************/
SearchBySet_::SearchBySet_( StatQuery* query, std::map<StatQuery::Key, SummaryIo>* rows )
{
    Search::query = query;
    Search::rows = rows;
}

/**************************************************************************************************************/
void
SearchBySet_::find( const StatRecord* begin, const StatRecord* end, time_t group, int hash )
{
    const StatRecord*   forward = begin;
    const StatRecord*   reverse;
    uint32_t            key[ 2 ];

    for ( const auto& n : query->networks() )
    {
        if ( (n.front() >> (32-PREFIX_HASH_LENGTH)) != hash )
        {
            continue;
        }

        key[ 0 ] = htonl( n.front() );
        key[ 1 ] = htonl( n.mask() );
        reverse = static_cast<const StatRecord*>( bsearch(key, forward, end-forward, sizeof(*end), compareNetwork_) );

        if ( reverse )
        {
            forward = reverse--;

            while ( reverse >= begin && !((reverse->beStat.ip&key[1]) ^ key[0]) )
            {
                extract( reverse--, group );
            }

            while ( forward < end && !((forward->beStat.ip&key[1]) ^ key[0]) )
            {
                extract( forward++, group );
            }
        }
    }
}
