#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 uint32_t PTR_LEN = sizeof(char*);
const uint32_t BAILOUT_MAX = 600; // 30s in 50ms increments
const uint32_t FILES_PER_WORK_UNIT = 40;

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;)
    
    // Start worker threads
    OUT(std::cout << "Starting worker threads..." << std::flush;)
    
    uint32_t num_threads = m_num_cores;
    for (uint32_t idx = 0; idx < num_threads; ++idx)
    {
        // Spawn worker thread and add to list for bookkeeping
        std::thread* worker = new std::thread(&PhotoCompressArchiver::fork_worker, 
                                              this, idx);
        m_worker_threads.push_back(worker);        
    }
    OUT(std::cout << "DONE\n" << std::endl;)
    
    // Join the file finder thread
    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 << " files" << std::endl;)
    
    if (! m_decompress)
    {
        OUT(std::cout << "  Total uncompressed bytes: " 
                      << m_total_uncompressed_size << "\n" << std::endl;)
    }
    
    
    // Build work unit queue
    uint32_t input_list_idx = 0;
    uint32_t queue_size = 0;
    while (input_list_idx < m_file_list.size())
    {
        // Safely read the size of the queue
        {
            std::lock_guard<std::mutex> lock(m_work_unit_mtx);
            queue_size = m_work_unit_queue.size();
        }
        
        // Wait until the queue has decreased in the size before proceeding,
        // otherwise we could chew up a bunch of memory
        if (queue_size > num_threads)
        {
             std::this_thread::sleep_for(std::chrono::milliseconds(250));
             continue;
        }
        
        WorkUnit* work_unit = new WorkUnit();
                
        work_unit->exec_argv.push_back(&PROG_NAME[2]);
        work_unit->exec_argv.push_back(PROG_OPT1);
        work_unit->exec_argv.push_back(PROG_OPT2);
        work_unit->exec_argv.push_back(PROG_OPT3);
    
        uint32_t argv_size = sizeof(PROG_NAME) + sizeof(PROG_OPT1) + 
                             sizeof(PROG_OPT2) + sizeof(PROG_OPT3) +  ARGV_EXTRA;
        
        uint32_t local_count = 0;
        while ((input_list_idx < m_file_list.size()) && 
               (local_count < FILES_PER_WORK_UNIT))
        {
            bfs::path& file_path = m_file_list[input_list_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
            work_unit->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");
            work_unit->output_filenames.push_back(filename);            
            
            // Increment our index to the next file
            ++input_list_idx;
            ++local_count;
        }
        work_unit->exec_argv.push_back(nullptr);
        
        
        // Add the create work unit to the queue; and notify the worker threads
        // that it has been handed off
        {
            std::lock_guard<std::mutex> lock(m_work_unit_mtx);
            m_work_unit_queue.push_back(work_unit);
        }
        m_work_unit_cv.notify_all();
    }
    
    // Add null terminate indicator for each thread
    {
        std::lock_guard<std::mutex> lock(m_work_unit_mtx);
        for (uint32_t idx = 0; idx < num_threads; ++idx)
        {

            m_work_unit_queue.push_back(nullptr);
        }
    }
    m_work_unit_cv.notify_all();
    
    // 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)
{
    // Loop vars
    int status;
    uint32_t work_unit_items = 0;
    uint32_t bailout_ctr = 0;
    uint32_t output_idx = 1;
    int stat_rc = 0;
    struct stat statbuf;
    uint32_t files_processed = 0;
    double complete_percent = 0.0;
    
    while (true)
    {
        
        std::unique_lock<std::mutex> mux_lock(m_work_unit_mtx);
        while (m_work_unit_queue.empty())
        {
            m_work_unit_cv.wait_for(mux_lock, std::chrono::milliseconds(20));
        }

        // While holding the mutex grab the next chunk off of the add chunk
        // queue
        WorkUnit* work_unit = m_work_unit_queue.front();
        m_work_unit_queue.pop_front();
        mux_lock.unlock();
        
        if (work_unit == nullptr)
        {
            break;
        }
        
        work_unit_items = work_unit->output_filenames.size();
        
#if 0
        // DEBUGGING
        uint32_t idx = 4;
        OUT(
        for (auto outfile : work_unit->output_filenames)
        {
            std::cout << "T[" << tid << "] output file: " << *outfile 
                          << "\n     input file:  " 
                          << static_cast<const char*>(work_unit->exec_argv[idx++]) 
                          << std::endl;
        }
        )
        continue;
#endif 
        
        // 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 **)work_unit->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.
        bailout_ctr = 0;
        output_idx = 1;
        stat_rc = 0;
        files_processed = 0;
        complete_percent = 0.0;
        while (output_idx < work_unit_items)
        {
            // Poll, waiting for output file X + 1 (aka "z") to exist
            bailout_ctr = 0;
            const char* filename_z = work_unit->output_filenames[output_idx]->c_str();
            while ((stat(filename_z, &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: " 
                              << *work_unit->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 (aka "y") which should exist
            const char* filename_y = work_unit->output_filenames[output_idx - 1]->c_str();
            stat_rc = stat(filename_y, &statbuf);
            if (stat_rc)
            {
                OUT(std::cerr << "T[" << tid << "] " << BOLD << RED << "ERROR " 
                              << ENDC << "while stating: " << filename_y 
                              << 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(work_unit->exec_argv[4 + output_idx - 1]);
            }
            
            // Calculate the local percent done
            complete_percent = (double)files_processed / work_unit_items;
            complete_percent *= 100;
            
            OUT(std::cout << "T[" << tid << "] " << std::setw(6)
                          << std::fixed << std::setprecision(2) 
                          << get_global_percent_done()
                          << "%  -- (" << std::setw(2) << files_processed << "/"
                          << work_unit_items << " -- " << std::setw(6)
                          << complete_percent 
                          << "%)    file: " << filename_y
                          << " -- " << 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
        const char* filename_y = work_unit->output_filenames[output_idx - 1]->c_str();
        stat_rc = stat(filename_y, &statbuf);
        if (stat_rc)
        {
            OUT(std::cerr << "T[" << tid << "] " << BOLD << RED << "ERROR " 
                          << ENDC << "while stating: " << filename_y << 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(work_unit->exec_argv[4 + output_idx - 1]);
        }
        
        // Calculate the local percent done
        complete_percent = (double)files_processed / work_unit_items;
        complete_percent *= 100;
        
        OUT(std::cout << "T[" << tid << "] " << std::setw(6) << std::fixed 
                      << std::setprecision(2) << get_global_percent_done() 
                      << "%  -- (" << std::setw(2) << files_processed 
                      << "/" << work_unit_items << " -- " << std::setw(5) 
                      << complete_percent << "%)    file: " << filename_y
                      << " -- " << statbuf.st_size << std::endl;)
        
        // Free up the output filename list
        for (auto output_filename : work_unit->output_filenames)
        {
            delete output_filename;
        }
        delete work_unit;
    }
    
    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;
}