// stat.cpp

/**
*    Copyright (C) 2008 10gen Inc.
*
*    This program is free software: you can redistribute it and/or  modify
*    it under the terms of the GNU Affero General Public License, version 3,
*    as published by the Free Software Foundation.
*
*    This program 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 Affero General Public License for more details.
*
*    You should have received a copy of the GNU Affero General Public License
*    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include "pch.h"
#include "client/dbclient.h"
#include "db/json.h"
#include "../util/httpclient.h"
#include "../util/text.h"

#include "tool.h"

#include <fstream>
#include <iostream>

#include <boost/program_options.hpp>

namespace po = boost::program_options;

namespace mongo {
    
    class Stat : public Tool {
    public:

        Stat() : Tool( "stat" , REMOTE_SERVER , "admin" ){
            _sleep = 1;
            _http = false;
            _many = false;
            
            add_hidden_options()
                ( "sleep" , po::value<int>() , "time to sleep between calls" )
                ;
            add_options()
                ("noheaders", "don't output column names")
                ("rowcount,n", po::value<int>()->default_value(0), "number of stats lines to print (0 for indefinite)")
                ("http", "use http instead of raw db connection")
                ("discover" , "discover nodes and display stats for all" )
                ("all" , "all optional fields" )
                ;

            addPositionArg( "sleep" , 1 );

            _autoreconnect = true;
        }

        virtual void printExtraHelp( ostream & out ){
            out << "usage: " << _name << " [options] [sleep time]" << endl;
            out << "sleep time: time to wait (in seconds) between calls" << endl;
        }

        virtual void printExtraHelpAfter( ostream & out ){
            out << "\n";
            out << " Fields\n";
            out << "   inserts \t- # of inserts per second\n";
            out << "   query   \t- # of queries per second\n";
            out << "   update  \t- # of updates per second\n";
            out << "   delete  \t- # of deletes per second\n";
            out << "   getmore \t- # of get mores (cursor batch) per second\n";
            out << "   command \t- # of commands per second\n";
            out << "   flushes \t- # of fsync flushes per second\n";
            out << "   mapped    \t- amount of data mmaped (total data size) megabytes\n";
            out << "   visze     \t- virtual size of process in megabytes\n";
            out << "   res       \t- resident size of process in megabytes\n";
            out << "   faults  \t- # of pages faults per sec (linux only)\n";
            out << "   locked    \t- percent of time in global write lock\n";
            out << "   idx miss  \t- percent of btree page misses (sampled)\n";
            out << "   qr|qw     \t- queue lengths for clients waiting (read|write)\n";
            out << "   ar|aw     \t- active clients (read|write)\n";
            out << "   netIn     \t- network traffic in - bits\n";
            out << "   netOut     \t- network traffic out - bits\n";
            out << "   conn      \t- number of open connections\n";
        }

        
        BSONObj stats(){
            if ( _http ){
                HttpClient c;
                HttpClient::Result r;
                
                string url;
                {
                    stringstream ss;
                    ss << "http://" << _host;
                    if ( _host.find( ":" ) == string::npos )
                        ss << ":28017";
                    ss << "/_status";
                    url = ss.str();
                }

                if ( c.get( url , &r ) != 200 ){
                    cout << "error (http): " << r.getEntireResponse() << endl;
                    return BSONObj();
                }
                
                BSONObj x = fromjson( r.getBody() );
                BSONElement e = x["serverStatus"];
                if ( e.type() != Object ){
                    cout << "BROKEN: " << x << endl;
                    return BSONObj();
                }
                return e.embeddedObjectUserCheck();
            }
            BSONObj out;
            if ( ! conn().simpleCommand( _db , &out , "serverStatus" ) ){
                cout << "error: " << out << endl;
                return BSONObj();
            }
            return out.getOwned();
        }

        double diff( const string& name , const BSONObj& a , const BSONObj& b ){
            BSONElement x = a.getFieldDotted( name.c_str() );
            BSONElement y = b.getFieldDotted( name.c_str() );
            if ( ! x.isNumber() || ! y.isNumber() )
                return -1;
            return ( y.number() - x.number() ) / _sleep;
        }
        
        double percent( const char * outof , const char * val , const BSONObj& a , const BSONObj& b ){
            double x = ( b.getFieldDotted( val ).number() - a.getFieldDotted( val ).number() );
            double y = ( b.getFieldDotted( outof ).number() - a.getFieldDotted( outof ).number() );
            if ( y == 0 )
                return 0;
            double p = x / y;
            p = (double)((int)(p * 1000)) / 10;
            return p;
        }

        template<typename T>
        void _append( BSONObjBuilder& result , const string& name , unsigned width , const T& t ){
            if ( name.size() > width )
                width = name.size();
            result.append( name , BSON( "width" << (int)width << "data" << t ) );
        }
        
        void _appendMem( BSONObjBuilder& result , const string& name , unsigned width , double sz ){
            string unit = "m";
            if ( sz > 1024 ){
                unit = "g";
                sz /= 1024;
            }

            if ( sz > 1024 ){
                string s = str::stream() << (int)sz << unit;
                _append( result , name , width , s );
                return;
            }

            stringstream ss;
            ss << setprecision(3) << sz << unit;
            _append( result , name , width , ss.str() );
        }
        
        void _appendNet( BSONObjBuilder& result , const string& name , double diff ){
            // I think 1000 is correct for megabit, but I've seen conflicting things (ERH 11/2010)
            const double div = 1000;
            
            string unit = "b";

            if ( diff >= div ){
                unit = "k";
                diff /= div;
            }
            
            if ( diff >= div ){
                unit = "m";
                diff /= div;
            }

            if ( diff >= div ){
                unit = "g";
                diff /= div;
            }

            string out = str::stream() << (int)diff << unit;
            _append( result , name , 6 , out );
        }

        /**
         * BSON( <field> -> BSON( width : ### , data : XXX ) )
         */
        BSONObj doRow( const BSONObj& a , const BSONObj& b ){
            BSONObjBuilder result;

            if ( b["opcounters"].type() == Object ){
                BSONObj ax = a["opcounters"].embeddedObject();
                BSONObj bx = b["opcounters"].embeddedObject();
                BSONObjIterator i( bx );
                while ( i.more() ){
                    BSONElement e = i.next();
                    _append( result , e.fieldName() , 6 , (int)diff( e.fieldName() , ax , bx ) );
                }
            }
            
            if ( b["backgroundFlushing"].type() == Object ){
                BSONObj ax = a["backgroundFlushing"].embeddedObject();
                BSONObj bx = b["backgroundFlushing"].embeddedObject();
                _append( result , "flushes" , 6 , (int)diff( "flushes" , ax , bx ) );
            }

            if ( b.getFieldDotted("mem.supported").trueValue() ){
                BSONObj bx = b["mem"].embeddedObject();
                BSONObjIterator i( bx );
                _appendMem( result , "mapped" , 6 , bx["mapped"].numberInt() );
                _appendMem( result , "vsize" , 6 , bx["virtual"].numberInt() );
                _appendMem( result , "res" , 6 , bx["resident"].numberInt() );

                if ( _all )
                    _appendMem( result , "non-mapped" , 6 , bx["virtual"].numberInt() - bx["mapped"].numberInt() );
            }

            if ( b["extra_info"].type() == Object ){
                BSONObj ax = a["extra_info"].embeddedObject();
                BSONObj bx = b["extra_info"].embeddedObject();
                if ( ax["page_faults"].type() || ax["page_faults"].type() )
                    _append( result , "faults" , 6 , (int)diff( "page_faults" , ax , bx ) );
            }
            
            _append( result , "locked %" , 8 , percent( "globalLock.totalTime" , "globalLock.lockTime" , a , b ) );
            _append( result , "idx miss %" , 8 , percent( "indexCounters.btree.accesses" , "indexCounters.btree.misses" , a , b ) );

            if ( b.getFieldDotted( "globalLock.currentQueue" ).type() == Object ){
                int r = b.getFieldDotted( "globalLock.currentQueue.readers" ).numberInt();
                int w = b.getFieldDotted( "globalLock.currentQueue.writers" ).numberInt();
                stringstream temp;
                temp << r << "|" << w;
                _append( result , "qr|qw" , 9 , temp.str() );
            }

            if ( b.getFieldDotted( "globalLock.activeClients" ).type() == Object ){
                int r = b.getFieldDotted( "globalLock.activeClients.readers" ).numberInt();
                int w = b.getFieldDotted( "globalLock.activeClients.writers" ).numberInt();
                stringstream temp;
                temp << r << "|" << w;
                _append( result , "ar|aw" , 7 , temp.str() );
            }
            
            if ( b["network"].isABSONObj() ){
                BSONObj ax = a["network"].embeddedObject();
                BSONObj bx = b["network"].embeddedObject();
                _appendNet( result , "netIn" , diff( "bytesIn" , ax , bx ) );
                _appendNet( result , "netOut" , diff( "bytesOut" , ax , bx ) );
            }

            _append( result , "conn" , 5 , b.getFieldDotted( "connections.current" ).numberInt() );

            if ( b["repl"].type() == Object ){

                BSONObj x = b["repl"].embeddedObject();
                bool isReplSet = x["setName"].type() == String;

                stringstream ss;

                if ( isReplSet ){
                    string setName = x["setName"].String();
                    _append( result , "set" , setName.size() , setName );
                }

                if ( x["ismaster"].trueValue() )
                    ss << "M";
                else if ( x["secondary"].trueValue() )
                    ss << "SEC";
                else if ( x["isreplicaset"].trueValue() )
                    ss << "REC";
                else if ( isReplSet )
                    ss << "UNK";
                else 
                    ss << "SLV";
                
                _append( result , "repl" , 4 , ss.str() );
                
            }
            else if ( b["shardCursorType"].type() == Object ){
                // is a mongos
                // TODO: should have a better check
                _append( result , "repl" , 4 , "RTR" );
            }

            {
                struct tm t;
                time_t_to_Struct( time(0), &t , true );
                stringstream temp;
                temp << setfill('0') << setw(2) << t.tm_hour 
                     << ":" 
                     << setfill('0') << setw(2) << t.tm_min
                     << ":" 
                     << setfill('0') << setw(2) << t.tm_sec;
                _append( result , "time" , 10 , temp.str() );
            }
            return result.obj();
        }
        
        virtual void preSetup(){
            if ( hasParam( "http" ) ){
                _http = true;
                _noconnection = true;
            }

            if ( hasParam( "host" ) && 
                 getParam( "host" ).find( ',' ) != string::npos ){
                _noconnection = true;
                _many = true;
            }

            if ( hasParam( "discover" ) ){
                _noconnection = true;
                _many = true;
            }
        }

        int run(){ 
            _sleep = getParam( "sleep" , _sleep );
            _all = hasParam( "all" );
            if ( _many )
                return runMany();
            return runNormal();
        }

        static void printHeaders( const BSONObj& o ){
            BSONObjIterator i(o);
            while ( i.more() ){
                BSONElement e = i.next();
                BSONObj x = e.Obj();
                cout << setw( x["width"].numberInt() ) << e.fieldName() << ' ';
            }
            cout << endl;            
        }

        static void printData( const BSONObj& o , const BSONObj& headers ){
                
            BSONObjIterator i(headers);
            while ( i.more() ){
                BSONElement e = i.next();
                BSONObj h = e.Obj();
                int w = h["width"].numberInt();
                
                BSONElement data;
                {
                    BSONElement temp = o[e.fieldName()];
                    if ( temp.isABSONObj() )
                        data = temp.Obj()["data"];
                }
                
                if ( data.type() == String )
                    cout << setw(w) << data.String();
                else if ( data.type() == NumberDouble )
                    cout << setw(w) << setprecision(3) << data.number();
                else if ( data.type() == NumberInt )
                    cout << setw(w) << data.numberInt();
                else if ( data.eoo() )
                    cout << setw(w) << "";
                else 
                    cout << setw(w) << "???";
                
                cout << ' ';
            }
            cout << endl; 
        }

        int runNormal(){
            bool showHeaders = ! hasParam( "noheaders" );
            int rowCount = getParam( "rowcount" , 0 );
            int rowNum = 0;

            BSONObj prev = stats();
            if ( prev.isEmpty() )
                return -1;

            while ( rowCount == 0 || rowNum < rowCount ){
                sleepsecs(_sleep);
                BSONObj now;
                try {
                    now = stats();
                }
                catch ( std::exception& e ){
                    cout << "can't get data: " << e.what() << endl;
                    continue;
                }

                if ( now.isEmpty() )
                    return -2;
                
                try {

                    BSONObj out = doRow( prev , now );

                    if ( showHeaders && rowNum % 10 == 0 ){
                        printHeaders( out );
                    }
                    
                    printData( out , out );

                }
                catch ( AssertionException& e ){
                    cout << "\nerror: " << e.what() << "\n"
                         << now
                         << endl;
                }
                
                prev = now;
                rowNum++;
            }
            return 0;
        }
        
        struct ServerState {
            ServerState() : lock( "Stat::ServerState" ){}
            string host;
            scoped_ptr<boost::thread> thr;
            
            mongo::mutex lock;
            
            BSONObj prev;
            BSONObj now;
            time_t lastUpdate;
            vector<BSONObj> shards;

            string error;
            bool mongos;
        };
            
        static void serverThread( shared_ptr<ServerState> state ){
            try {
                DBClientConnection conn( true );
                conn._logLevel = 1;
                string errmsg;
                if ( ! conn.connect( state->host , errmsg ) )
                    state->error = errmsg;
                
                long long cycleNumber = 0;

                while ( ++cycleNumber ){
                    try {
                        BSONObj out;
                        if ( conn.simpleCommand( "admin" , &out , "serverStatus" ) ){
                            scoped_lock lk( state->lock );
                            state->error = "";
                            state->lastUpdate = time(0);
                            state->prev = state->now;
                            state->now = out.getOwned();
                        }
                        else {
                            scoped_lock lk( state->lock );
                            state->error = "serverStatus failed";
                            state->lastUpdate = time(0);
                        }
                        
                        if ( out["shardCursorType"].type() == Object ){
                            state->mongos = true;
                            if ( cycleNumber % 10 == 1 ){
                                auto_ptr<DBClientCursor> c = conn.query( "config.shards" , BSONObj() );
                                vector<BSONObj> shards;
                                while ( c->more() ){
                                    shards.push_back( c->next().getOwned() );
                                }
                                scoped_lock lk( state->lock );
                                state->shards = shards;
                            }
                        }
                    }
                    catch ( std::exception& e ){
                        scoped_lock lk( state->lock );
                        state->error = e.what();
                    }
                    
                    sleepsecs( 1 );
                }
                
                
            }
            catch ( std::exception& e ){
                cout << "serverThread (" << state->host << ") fatal error : " << e.what() << endl;
            }
            catch ( ... ){
                cout << "serverThread (" << state->host << ") fatal error" << endl;
            }
        }

        typedef map<string,shared_ptr<ServerState> >  StateMap;

        bool _add( StateMap& threads , string host ){
            shared_ptr<ServerState>& state = threads[host];
            if ( state )
                return false;
            
            state.reset( new ServerState() );
            state->host = host;
            state->thr.reset( new boost::thread( boost::bind( serverThread , state ) ) );
            return true;
        }
        
        /**
         * @param hosts [ "a.foo.com" , "b.foo.com" ]
         */
        bool _addAll( StateMap& threads , const BSONObj& hosts ){
            BSONObjIterator i( hosts );
            bool added = false;
            while ( i.more() ){
                bool me = _add( threads , i.next().String() );
                added = added || me;
            }
            return added;
        }

        bool _discover( StateMap& threads , const string& host , const shared_ptr<ServerState>& ss ){
            
            BSONObj info = ss->now;

            bool found = false;
            
            if ( info["repl"].isABSONObj() ){
                BSONObj x = info["repl"].Obj();
                if ( x["hosts"].isABSONObj() )
                    if ( _addAll( threads , x["hosts"].Obj() ) )
                        found = true;
                if ( x["passives"].isABSONObj() )
                    if ( _addAll( threads , x["passives"].Obj() ) )
                        found = true;
            }
            
            if ( ss->mongos ){
                for ( unsigned i=0; i<ss->shards.size(); i++ ){
                    BSONObj x = ss->shards[i];

                    string errmsg;
                    ConnectionString cs = ConnectionString::parse( x["host"].String() , errmsg );
                    if ( errmsg.size() ){
                        cerr << errmsg << endl;
                        continue;
                    }
                    
                    vector<HostAndPort> v = cs.getServers();
                    for ( unsigned i=0; i<v.size(); i++ ){
                        if ( _add( threads , v[i].toString() ) )
                            found = true;
                    }
                }
            }
            
            return found;
        }
        
        int runMany(){
            StateMap threads;
            

            {
                string orig = getParam( "host" );
                if ( orig == "" )
                    orig = "localhost:27017";
                StringSplitter ss( orig.c_str() , "," );
                while ( ss.more() ){
                    string host = ss.next();
                    _add( threads , host );
                }
            }
            
            sleepsecs(1);
            
            int row = 0;
            bool discover = hasParam( "discover" );

            while ( 1 ){
                sleepsecs( _sleep );
                
                // collect data
                vector<Row> rows;
                for ( map<string,shared_ptr<ServerState> >::iterator i=threads.begin(); i!=threads.end(); ++i ){
                    scoped_lock lk( i->second->lock );
                    
                    if ( i->second->error.size() ){
                        rows.push_back( Row( i->first , i->second->error ) );
                    }
                    else if ( i->second->prev.isEmpty() || i->second->now.isEmpty() ){
                        rows.push_back( Row( i->first ) );
                    }
                    else {
                        BSONObj out = doRow( i->second->prev , i->second->now );
                        rows.push_back( Row( i->first , out ) );
                    }
                    
                    if ( discover && ! i->second->now.isEmpty() ){
                        if ( _discover( threads , i->first , i->second ) )
                            break;
                    }
                }
                
                // compute some stats
                unsigned longestHost = 0;
                BSONObj biggest;
                for ( unsigned i=0; i<rows.size(); i++ ){
                    if ( rows[i].host.size() > longestHost )
                        longestHost = rows[i].host.size();
                    if ( rows[i].data.nFields() > biggest.nFields() )
                        biggest = rows[i].data;
                }
                
                { // check for any headers not in biggest

                    // TODO: we put any new headers at end, 
                    //       ideally we would interleave

                    set<string> seen;
                    
                    BSONObjBuilder b;
                    
                    { // iterate biggest
                        BSONObjIterator i( biggest );
                        while ( i.more() ){
                            BSONElement e = i.next();
                            seen.insert( e.fieldName() );
                            b.append( e );
                        }
                    }
                    
                    // now do the rest
                    for ( unsigned j=0; j<rows.size(); j++ ){
                        BSONObjIterator i( rows[j].data );
                        while ( i.more() ){
                            BSONElement e = i.next();
                            if ( seen.count( e.fieldName() ) )
                                continue;
                            seen.insert( e.fieldName() );
                            b.append( e );
                        }

                    }

                    biggest = b.obj();
                    
                }
                
                // display data
                
                cout << endl;

                //    header
                if ( row++ % 5 == 0 && ! biggest.isEmpty() ){
                    cout << setw( longestHost ) << "" << "\t";
                    printHeaders( biggest );
                }
                
                //    rows
                for ( unsigned i=0; i<rows.size(); i++ ){
                    cout << setw( longestHost ) << rows[i].host << "\t";
                    if ( rows[i].err.size() )
                        cout << rows[i].err << endl;
                    else if ( rows[i].data.isEmpty() )
                        cout << "no data" << endl;
                    else 
                        printData( rows[i].data , biggest );
                }
                
            }

            return 0;
        }

        int _sleep;
        bool _http;
        bool _many;
        bool _all;

        struct Row {
            Row( string h , string e ){
                host = h;
                err = e;
            }
            
            Row( string h ){
                host = h;
            }

            Row( string h , BSONObj d ){
                host = h;
                data = d;
            }
            string host;
            string err;
            BSONObj data;
        };
    };

}

int main( int argc , char ** argv ) {
    mongo::Stat stat;
    return stat.main( argc , argv );
}