#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));
const uint32_t ARGV_EXTRA(4096);

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;
    
    
    // Remove this...
    std::cout << "argv size: " << sysconf(_SC_ARG_MAX) << 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;)
    
    // Error out of no input files were found
    if (m_num_input_files == 0)
    {
        OUT(std::cerr << BOLD << RED << "ERROR: " << ENDC << "No matching "
                      << "input files found. Please check the patch and try "
                      << "again. " << std::endl;)
        return -1;
    }
    
    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 = 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;)
    }
    
    return; // temporary for debugging
#endif

    uint32_t sublist_idx = 0;
    const uint32_t ptr_len = sizeof(char*);
    int status;
    
    // Vector of pointers that will become the char** argv passed to packjpg
    std::vector<const char*> exec_argv;
    
    // 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;
    
    while (sublist_idx < sublist_file_count)
    {
        // Pass in default program options for packjpg
        exec_argv.clear();
        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);
    
        uint32_t argv_size = sizeof(PROG_NAME) + sizeof(PROG_OPT1) + 
                             sizeof(PROG_OPT2) + sizeof(PROG_OPT3) +  ARGV_EXTRA;
        
        output_filenames.clear();
        
        uint32_t local_count = 0;
        while ((sublist_idx < sublist_file_count) && (local_count < 40))
        {
            bfs::path& file_path = file_sublist->at(sublist_idx);
            uint32_t path_len = file_path.string().size();
            
            // Check if adding this next file would make the argv too big; if
            // so bail out
            if (argv_size + path_len + ptr_len > ARGV_MAX_SIZE)
            {
                break;
            }
            
            // Add the file to the argv array
            exec_argv.push_back(file_path.string().c_str());
            argv_size += path_len + ptr_len;
            
             // Create new filename for expected 
            uint32_t end_pos = file_path.string().rfind(".");
            std::string* filename = new std::string(file_path.string().begin(), 
                                                    file_path.string().begin() 
                                                    + end_pos);
                
            // Add extension depending on if decompress mode is enabled or not
            (*filename) += (m_decompress ? ".jpg" : ".pjg");
            output_filenames.push_back(filename);
            
            
            // Increment our index to the next file
            ++sublist_idx;
            ++local_count;
        }
        
#if 0
        // DEBUGGING
        uint32_t idx = 4;
        OUT(
        for (auto outfile : output_filenames)
        {
            std::cout << "T[" << tid << "] output file: " << *outfile 
                          << "\n     input file:  " 
                          << static_cast<const char*>(exec_argv[idx++]) << std::endl;
        }
        )
        //~ continue;
#endif 
        
        // Add a nullptr to the end of the argv array as a terminator
        exec_argv.push_back(nullptr);

        // 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[0]);
            close(filedes[1]);
            
            // 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;)
                
                // The subprocess timed out... get its 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 << "T[" << tid << "] " << 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
        waitpid(pid, &status, 0);
        close(filedes[0]);
        close(filedes[1]);
        
        // 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 << "T[" << tid << "] " << 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;)
        
        // Free up the output filename list
        for (auto output_filename : output_filenames)
        {
            delete output_filename;
        }
    }
    
    OUT(std::cout << "T[" << tid << "] " << PURPLE << "complete ==> RC: " 
                  << status << ENDC << std::endl;)
}

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