#include <chrono>
#include <cstdio>
#include <iomanip>
#include <iostream>
#include <limits.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <unistd.h>

#include "PhotoCompressArchiver.hh"

#define OUT(x) {\
    std::lock_guard<std::mutex> lock(m_output_mutex);\
    x\
}

const char *PROG_NAME = "./packjpg";
const char *PROG_OPT1 = "-np";
const char *PROG_OPT2 = "-o";
const char *PROG_OPT3 = "-p";

const uint32_t ARGV_MAX_SIZE(sysconf(_SC_ARG_MAX));

PhotoCompressArchiver::PhotoCompressArchiver(
    const std::string& path, 
    uint32_t num_cores,
    bool decompress,
    bool keep_orig,
    boost::regex& filter_regex)
    : m_path(path),
      m_num_cores(num_cores),
      m_decompress(decompress),
      m_keep_orig(keep_orig),
      m_files_processed(0),
      m_num_input_files(0),
      m_total_uncompressed_size(0),
      m_total_compressed_size(0),
      p_finder_thread(nullptr),
      m_filter_empty(true),
      m_filter_regex(filter_regex)
{
    // This is a somewhat crude way to determining if the filter regex is empty
    m_filter_empty = std::string("").compare(filter_regex.str()) == 0;
    
    std::cout << "argv size: " << sysconf(_SC_ARG_MAX) << std::endl;
    std::cout << "appname size: " << sizeof(PROG_NAME) << std::endl;
}

PhotoCompressArchiver::~PhotoCompressArchiver()
{}

int PhotoCompressArchiver::execute()
{
    // Start the file finder thread
    OUT(std::cout << "Starting file finder..." << std::flush;);
    p_finder_thread = new std::thread(&PhotoCompressArchiver::find_files, this, 
                                    m_path, 
                                    (m_decompress ? PJG_REGEX : JPEG_REGEX));
    OUT(std::cout << "DONE" << std::endl;)
    p_finder_thread->join();
    OUT(std::cout << "File finder completed:" << std::endl;)
    
    OUT(std::cout << "  Found " << m_num_input_files << " images" << std::endl;)
    
    if (! m_decompress)
    {
        OUT(std::cout << "  Total uncompressed bytes: " 
                      << m_total_uncompressed_size << "\n" << std::endl;)
    }
    
    // Start worker threads
    OUT(std::cout << "Starting worker threads..." << std::flush;)
    
    // TODO: need to handle the case were number of input files is less than
    //       the number of threads
    uint32_t num_threads = 1;//m_num_cores;
    uint32_t sublist_len = m_num_input_files / num_threads;
 
    for (uint32_t idx = 0; idx < num_threads; ++idx)
    {
        // Subdivide the input file list (using the vector copy constructor 
        // and start/end iterators)
        decltype(m_file_list.begin()) start_iter;
        decltype(m_file_list.begin()) end_iter;
        if (idx == (num_threads - 1))
        {
            // Make sure the last sublist goes all the way to the end of the 
            // input file list (if number of images is not an even multiple of
            // number of threads)
            start_iter = m_file_list.begin() + (idx * sublist_len);
            end_iter   = m_file_list.end();
        }
        else
        {
            start_iter = m_file_list.begin() + (idx * sublist_len);
            end_iter   = m_file_list.begin() + ((idx + 1) * sublist_len);
        }
        
        auto sublist = new std::vector<bfs::path>(start_iter, end_iter);
        
        // Spawn worker thread and add to list for bookkeeping
        std::thread* worker = new std::thread(&PhotoCompressArchiver::fork_worker, 
                                              this, idx, sublist);
        m_worker_threads.push_back(worker);        
    }
    OUT(std::cout << "DONE\n" << std::endl;)
    
    // Clean up the finder thread
    delete p_finder_thread;
    p_finder_thread = nullptr;
    
    // Now join the worker threads
    for (auto worker : m_worker_threads)
    {
        worker->join();
        delete worker;
    }
    
    // Print out compression information if compressing
    if (! m_decompress)
    {
        OUT(std::cout << "\n  Total compressed bytes: " 
                      << m_total_compressed_size 
                      << std::endl;)
        
        // Calculate and display ratio
        double ratio = (double)m_total_compressed_size / m_total_uncompressed_size;
        OUT(std::cout << "  Compression ratio: " << std::fixed 
                      << std::setprecision(4) << (ratio * 100) << "%\n" 
                      << std::endl;)
    }
    
    return 0;
}

void PhotoCompressArchiver::find_files(
    const bfs::path& dir_path, 
    const boost::regex& file_regex)
{
    bfs::directory_iterator end_iter;
    for (bfs::directory_iterator dir_iter(dir_path); dir_iter != end_iter; ++dir_iter)
    {
        if (bfs::is_directory(dir_iter->status()))
        {
            // Found directory recurse
            find_files(dir_iter->path(), file_regex);
        }
        else if (boost::regex_match(dir_iter->path().filename().string(), file_regex))
        {
            // Found file match
            
            // If the filter is not empty, then check against it before proceeding
            if (! m_filter_empty)
            {
                if (boost::regex_match(dir_iter->path().filename().string(), 
                    m_filter_regex))
                {
                    // The filter regex matched, skip this file
                    continue;
                }
            }
            
            m_file_list.push_back(dir_iter->path());
            ++m_num_input_files;
            
            // Stat the file to get its size
            struct stat statbuf;
            int rc = stat(dir_iter->path().string().c_str(), &statbuf);
            if (rc)
            {
                OUT(std::cerr << BOLD << RED << "ERROR: " << ENDC 
                              << "unable to stat file " << dir_iter->path() 
                              << std::endl;)
            }     
            m_total_uncompressed_size += statbuf.st_size;
        }
    }
}

void PhotoCompressArchiver::fork_worker(
    uint32_t tid, 
    std::vector<bfs::path>* file_sublist)
{
    uint32_t sublist_file_count = file_sublist->size();
    uint32_t files_processed = 0;
    
#if 0
    // Print out the input files given to this worker
    OUT(std::cout << "T[" << tid << "] top of the morning to ya: " 
                  << sublist_file_count << std::endl;)

    for (auto& filepath : *file_sublist)
    {
        OUT(std::cout << "T[" << tid << "] file: " << filepath.string() 
                      << std::endl;)
    }
#endif

    //~ return; // temporary for debugging
        
    //~ uint32_t argv_size = file_sublist->size() + 5;
    //~ const char **exec_argv = new const char* [argv_size];
    
    // Build the argv array; pass in default program options
    std::vector<const char*> exec_argv;
    exec_argv.push_back(&PROG_NAME[2]);
    exec_argv.push_back(PROG_OPT1);
    exec_argv.push_back(PROG_OPT2);
    exec_argv.push_back(PROG_OPT3);
    
    // Build the argv array; pass in default program options
    //~ exec_argv[0] = &PROG_NAME[2]; // discard leading "./"
    //~ exec_argv[1] = PROG_OPT1;
    //~ exec_argv[2] = PROG_OPT2;
    //~ exec_argv[3] = PROG_OPT3;
    
    uint32_t argv_size = sizeof(PROG_NAME) + sizeof(PROG_OPT1) + 
                         sizeof(PROG_OPT2) + sizeof(PROG_OPT3) +  4096;
    
    // Add the sublist files to the argv array
    uint32_t idx = 0;
    for (auto& image_path : *file_sublist)
    {
        uint32_t image_path_len = image_path.string().size();
        if (argv_size + image_path_len < ARGV_MAX_SIZE)
        {
            exec_argv.push_back(image_path.string().c_str());
            argv_size += image_path_len + sizeof(char*);
            ++idx;
        }
        else
        {
            break;
        }
        
        //~ exec_argv[idx++] = image_path.string().c_str();
    }
    //~ exec_argv[argv_size - 1] = nullptr;
    exec_argv.push_back(nullptr);
    
    std::cout << "IDX: " << idx << std::endl;
    
    
    // Build a list of the output files names for use in output file monitoring
    // below. Output file names have ".jpg" or ".jpeg" replaced with ".pjg"
    // in compression mode and vice versa in decompression mode.
    std::vector<std::string*> output_filenames;
    for (auto& filepath : *file_sublist)
    {
        // Create new filename 
        uint32_t end_pos = filepath.string().find(".");
        std::string* filename = new std::string(filepath.string().begin(), 
                                                filepath.string().begin() + end_pos);
        
        // Add extension depending on if decompress mode is enabled
        (*filename) += (m_decompress ? ".jpg" : ".pjg");
        
        output_filenames.push_back(filename);
    }
    
    // Create a pipe to hold stdout from child process
    int filedes[2];
    if (pipe(filedes) == -1) 
    {
        perror("pipe");
        exit(1);
    }
    
    // Fork off the childprocess
    //~ pid_t parent = getpid();
    pid_t pid = vfork();

    if (pid == -1)
    {
        // error, failed to fork()
        OUT(std::cout << "T[" << tid << "] fork failed!" << std::endl;)
    } 
    else if (pid == 0)
    {
        // This is the child process...
        // Set child process's stdout to the pipe entry
        //~ while ((dup2(filedes[1], STDOUT_FILENO) == -1) && (errno == EINTR)) {}
        //~ close(filedes[1]);
        //~ close(filedes[0]);
        
        // Child after fork
        int rc = execv(PROG_NAME, (char **)exec_argv.data());
        if (rc)
        {
            std::cerr << "execv failed: " << errno << std::endl;
        }
        _exit(EXIT_FAILURE);   // exec never returns
    }
    
    // This is the parent process...
    
    // This is a little silly... but it works. We know the order in which the
    // output files will be created so we wait until output files X + 1 exists
    // (i.e. we can stat() it) and then process file X.
    const uint32_t BAILOUT_MAX = 600; // 30s in 50ms increments
    uint32_t bailout_ctr = 0;
    uint32_t output_idx = 1;
    int stat_rc = 0;
    struct stat statbuf;
    double complete_percent = 0.0;
    while (output_idx < output_filenames.size())
    {
        // Poll, waiting for output file X + 1 to exist
        bailout_ctr = 0;
        while ((stat(output_filenames[output_idx]->c_str(), &statbuf) != 0) &&
               (++bailout_ctr < BAILOUT_MAX))
        {
            std::this_thread::sleep_for(std::chrono::milliseconds(50));
        }
        
        // Check if timeout has occurred
        if (bailout_ctr == BAILOUT_MAX) 
        {
            OUT(std::cerr << "T[" << tid << "] " << BOLD << RED << "ERROR: " 
                          << ENDC << "timed out while waiting "
                          << "for output file: " 
                          << *output_filenames[output_idx] << "; aborting"
                          << std::endl;)
                          
            int status;
            waitpid(pid, &status, 0);
            
            OUT(std::cerr << "T[" << tid << "] " << BOLD << RED << "ERROR: " 
                          << ENDC << PURPLE << " RC => " << status << ENDC 
                          << std::endl;)
            
            return;
        }
                
        // Now stat output file X which should exist
        stat_rc = stat(output_filenames[output_idx - 1]->c_str(), &statbuf);
        if (stat_rc)
        {
            OUT(std::cerr << BOLD << RED << "ERROR " << ENDC 
                          << "while stating: " 
                          << *output_filenames[output_idx - 1] << std::endl;)
        }
        
        // Increment counts
        m_total_compressed_size += statbuf.st_size;
        ++m_files_processed;
        ++files_processed;
        
        // Unless keeping original files, remove the source file
        if (! m_keep_orig)
        {
            std::remove(file_sublist->at(output_idx - 1).string().c_str());
        }
        
        // Calculate the local percent done
        complete_percent = (double)files_processed / sublist_file_count;
        complete_percent *= 100;
        
        OUT(std::cout << "T[" << tid << "] " << std::setw(6)
                      << std::fixed << std::setprecision(2) 
                      << get_global_percent_done()
                      << "%  -- (" << files_processed << "/"
                      << sublist_file_count << " -- " << std::setw(6)
                      << complete_percent 
                      << "%)    file: " << *output_filenames[output_idx - 1] 
                      << " -- " << statbuf.st_size << std::endl;)
        
        ++output_idx;
    }
    
    // Wait for forked child process to terminate
    int status;
    waitpid(pid, &status, 0);
    
    // Now that the forked process has completed, we can handle the last output
    // file
    stat_rc = stat(output_filenames[output_idx - 1]->c_str(), &statbuf);
    if (stat_rc)
    {
        OUT(std::cerr << BOLD << RED << "ERROR " << ENDC << "while stating: " 
                      << *output_filenames[output_idx - 1] << std::endl;)
    }
    
    // Increment counts
    m_total_compressed_size += statbuf.st_size;
    ++m_files_processed;
    ++files_processed;
    
    // Unless keeping original files, remove the source file
    if (! m_keep_orig)
    {
        std::remove(file_sublist->at(output_idx - 1).string().c_str());
    }
    
    // Calculate the local percent done
    complete_percent = (double)files_processed / sublist_file_count;
    complete_percent *= 100;
    
    OUT(std::cout << "T[" << tid << "] " << std::setw(6) << std::fixed 
                  << std::setprecision(2) << get_global_percent_done() 
                  << "%  -- (" << files_processed << "/" << sublist_file_count 
                  << " -- " << std::setw(5) << complete_percent << "%)    file: " 
                  << *output_filenames[output_idx - 1] << " -- " 
                  << statbuf.st_size << std::endl;)
    
    OUT(std::cout << "T[" << tid << "] "<< PURPLE << "complete ==> RC: " 
                  << status << ENDC << std::endl;)
    
    // Free up the argv array and output filename list
    //~ delete[] exec_argv;
    for (auto output_filename : output_filenames)
    {
        delete output_filename;
    }
}

double PhotoCompressArchiver::get_global_percent_done()
{
    return ((double)m_files_processed / m_num_input_files) * 100;
}