#include <windows.h>
#include "st.h"
#include <sstream>
#include <iomanip>
#include <cstring>

using namespace std;

#ifdef NO_MAIN
extern "C" {
extern int df_main(int argc, char ** argv);
}
int df_main(int argc, char ** argv) 
#else
int main(int argc, char ** argv) 
#endif
{

    DWORD drives = GetLogicalDrives();
    string dr;
    ULARGE_INTEGER free_bytes;
    ULARGE_INTEGER total_bytes;
    ULARGE_INTEGER call_bytes;
    long long div = 1024 * 1024 * 1024;
    string size = "GB";
    int i;

    for(i = 1; i < argc; i++) {
       if(strcmp(argv[i], "-b") == 0) {
          div = 1;
          size = "B";
       }
       else if(strcmp(argv[i], "-k") == 0) {
          div = 1024;
          size = "KB";
       }
       else if(strcmp(argv[i], "-m") == 0) {
          div = 1024 * 1024;
          size = "MB";
       }
       else if(strcmp(argv[i], "-g") == 0) {
          div = 1024 * 1024 * 1024;
          size = "GB";
       }
    }

    if (drives == 0) {
        printf("Error getting logical drives.\n");
        return 1;
    }

    for (char drive = 'A'; drive <= 'Z'; drive++) {
        if (drives & 1) {
            stringstream str;
            dr = drive + string(":\\");
            str << setw(5) << left << dr;
            if(GetDiskFreeSpaceEx(dr.c_str(), &call_bytes, &total_bytes, &free_bytes)) {
               str << setw(5) << right << "free: " << setw(5) << right << (free_bytes.QuadPart/div) << " " << size  
                   << setw(5) << right << " total: " << setw(5) << right << (total_bytes.QuadPart/div) << " " << size << endl;
            }
            else {
               str << endl;
            }
            cout << str.str();
        }
        drives >>= 1; 
    }

    return 0;
}
