#include <iostream>
#include <list>
#include <string>
#include <map>
#include <cstring>

#include <sys/types.h>
#include <sys/stat.h>
#include <dirent.h>
#include <unistd.h>

#include <boost/cstdint.hpp>

struct file_stats {
    size_t count;
    boost::int64_t size;
    file_stats () : count (0), size (0) {}
    friend std::ostream& operator<< (std::ostream& os, const file_stats& s) {
        os << s.count << " files, ";
        boost::int64_t divisor = 1 << 10;
        static char units[] = "kMGTYFZ";
        const char* unit = units;
        while (s.size > divisor << 10 && *(unit + 1))
            {
            divisor <<= 10;
            ++unit;
            }
        return os << (double) s.size / divisor << ' ' << *unit << "iB";
    }
};

typedef std::map<std::string, file_stats> stat_map;

void
finder (const std::string& root, stat_map& stats, bool recursive)
{
    DIR* dir = opendir (root.c_str ());

    if (!dir)
        return;

    struct dirent* entry;
    struct stat st;

    while ((entry = readdir (dir)) != 0)
        {
        if (strcmp (entry->d_name, ".") == 0 ||
            strcmp (entry->d_name, "..") == 0)
            continue;

        std::string full_path = root + '/' + entry->d_name;

        if (::stat (full_path.c_str(), &st) != 0)
            continue;

        if (S_ISDIR (st.st_mode) && recursive)
            finder (full_path, stats, recursive);
        else if (S_ISREG (st.st_mode))
            {
            const char* dot = strrchr (entry->d_name, '.');

            if (dot != 0)
                {
                std::string extension (dot);
                stat_map::iterator stat = stats.find (extension);

                if (stat != stats.end())
                    {
                    stat->second.count++;
                    stat->second.size += st.st_size;
                    }
                }
            }
        }

    closedir (dir);
}

int
main (int argc, char** argv)
{
    bool recursive = false;
    std::list<std::string> roots;
    stat_map stats;
    struct stat st;

    for (int i = 1; i < argc; ++i)
        {
        if (strcmp (argv[i], "-R") == 0)
            recursive = true;
        else if (::stat (argv[i], &st) == 0 && S_ISDIR (st.st_mode))
            roots.push_back (argv[i]);
        else
            stats.insert (std::make_pair (std::string (argv[i]),
                                          file_stats()));
        }

    if (stats.empty () || roots.empty ())
        {
        std::cerr << "Usage: " << argv[0]
                  << " [-R] directory [directory ...] " 
                  << "extension [extension ...]\n";
        exit (1);
        }

    for (std::list<std::string>::const_iterator iter = roots.begin(),
         end = roots.end(); iter != end; ++iter)
        {
        finder (*iter, stats, recursive);
        }

    for (stat_map::const_iterator iter = stats.begin(), end = stats.end ();
         iter != end; ++iter)
        std::cout << iter->first << ": " << iter->second << '\n';

    return 0;
}