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

#include <packjpg.h>

#include "PhotoCompressArchiver.hh"

#define ERROR     std::cerr << BOLD << RED << "ERROR: " << ENDC 

#define WRKR_OUT_REG(msg) {                                                   \
    std::lock_guard<std::mutex> lock(m_output_mutex);                         \
    std::cout << "T[" << std::setw(2) << tid << "] " << msg << std::endl;     \
}

#define WRKR_OUT_ERR(msg) {                                                   \
    std::lock_guard<std::mutex> lock(m_output_mutex);                         \
    std::cerr << "T[" << std::setw(2) << tid << "] " << BOLD << RED           \
              << "ERROR: " << ENDC << msg << std::endl;                       \
}

// Helper functions
std::string _format_num_bytes(uint64_t num_bytes)
{
    const uint64_t tera = (1024 * 1024 * 1024 * 1024L);
    const uint64_t giga = (1024 * 1024 * 1024L);
    const uint64_t mega = (1024 * 1024L);
    const uint64_t kilo = 1024L;
    
    double fractional; 
    std::ostringstream output;
    
    if (num_bytes >= tera)
    {
        // Tera
        fractional = ((double)num_bytes) / tera;
        output << std::fixed << std::setprecision(2) << fractional;
        output << " TB";
    }
    else if (num_bytes >= giga)
    {
        // Giga
        fractional = ((double)num_bytes) / giga;
        output << std::fixed << std::setprecision(2) << fractional;
        output << " GB";
    }
    else if (num_bytes >= mega)
    {
        // Mega
        fractional = ((double)num_bytes) / mega;
        output << std::fixed << std::setprecision(2) << fractional;
        output << " MB";
    }
    else if (num_bytes >= kilo)
    {
        // Kilo
        fractional = ((double)num_bytes) / kilo;
        output << std::fixed << std::setprecision(2) << fractional;
        output << " KB";
    }
    else
    {
        // Bytes
        output << num_bytes << " b";
    }
    
    return output.str();
}

PhotoCompressArchiver::PhotoCompressArchiver(
    const std::string& path, 
    uint32_t num_threads,
    bool decompress,
    bool delete_orig,
    boost::regex& filter_regex)
    : m_path(path),
      m_num_threads(num_threads),
      m_decompress(decompress),
      m_delete_orig(delete_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;
}

PhotoCompressArchiver::~PhotoCompressArchiver()
{}

int PhotoCompressArchiver::execute()
{
    // Grab start time
    auto start = std::chrono::system_clock::now();
    
    // Start the file finder thread
    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));
    std::cout << "DONE" << std::endl;
    
    // Start worker threads
    std::cout << "Starting worker threads..." << std::flush;
    for (uint32_t idx = 0; idx < m_num_threads; ++idx)
    {
        // Spawn worker thread and add to list for bookkeeping
        std::thread* worker = new std::thread(&PhotoCompressArchiver::worker_thread_body, 
                                              this, idx);
        m_worker_threads.push_back(worker);        
    }
    std::cout << "DONE\n" << std::endl;
    
    // Join the file finder thread
    p_finder_thread->join();
    std::cout << "File finder completed:" << std::endl;
    
    // Error out of no input files were found
    if (m_num_input_files == 0)
    {
        ERROR << "No matching input files found. Please check the patch and "
              << "try again." << std::endl;
        return -1;
    }
    
    std::cout << "  Found " << m_num_input_files << " files" << std::endl;
    
    if (! m_decompress)
    {
        std::cout << "  Total uncompressed bytes: " 
                  << _format_num_bytes(m_total_uncompressed_size) << "\n" 
                  << std::endl;
    }
    
    // Add terminator for each worker thread
    for (uint32_t idx = 0; idx < m_num_threads; ++idx)
    {
        m_file_list.push_back(nullptr);
    }
    
    // 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;
    }
    
    // Grab start time and calculate duration
    auto end = std::chrono::system_clock::now();
    auto diff = end - start;
    double duration = std::chrono::duration<double>(diff).count();
    
    // Print out message if any errors were found
    if (m_errors.size() > 0)
    {
        std::cout << BOLD << RED << "\n\nNOTE: errors were encountered "
                  << "processing the following " << m_errors.size() 
                  << " files:\n" << ENDC << std::endl;
            
        WorkUnit* work_unit = nullptr;
        uint32_t size = m_errors.size();
        for (uint32_t idx = 0; idx < size; ++idx)
        {
            work_unit = m_errors.pop_front();
            std::cout << "    " << work_unit->m_path.string() << std::endl;
            delete work_unit;
        }
        std::cout << std::endl;
    }
    
    // Print out compression information if compressing
    if (! m_decompress)
    {
        std::cout << "\n  Total uncompressed bytes: " 
                  << _format_num_bytes(m_total_uncompressed_size) << std::endl;
        std::cout << "  Total compressed bytes:   " 
                  << _format_num_bytes(m_total_compressed_size) << std::endl;
        
        // Calculate and display ratio
        double ratio = (double)m_total_compressed_size / m_total_uncompressed_size;
        std::cout << "  Compression ratio:        " << std::fixed 
                  << std::setprecision(2) << (ratio * 100) << "%" << std::endl;
    }
    else
    {
        std::cout << std::endl;
    }
    
    // Display execution duration
    std::cout << "  Total time:               " << duration << " s\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) &&
                 bfs::is_regular_file(dir_iter->path()))
        {
            // 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;
                }
            }
            
            // Create a WorkUnit and add it to the queue
            WorkUnit* work_unit = new WorkUnit(dir_iter->path());
            ++m_num_input_files;
            m_total_uncompressed_size += work_unit->m_file_size;
            m_file_list.push_back(work_unit);
        }
    }
}

void PhotoCompressArchiver::worker_thread_body(uint32_t tid)
{
    WorkUnit* work_unit = nullptr;
    auto file_buffer = new std::vector<unsigned char>(4 * 1024 * 1024);
    
    // Setup packJPG instance for this worker
    packJPG* instance = new packJPG();
    
    // (De/)compress worker variables
    unsigned char* out_buffer = nullptr;
    uint32_t out_size = 0;
    char message[MSG_SIZE];
    
    while (true)
    {
        // Reset worker vars
        out_size = 0;
        std::memset(message, 0, MSG_SIZE);
        
        // Blocking wait for next work unit
        work_unit = m_file_list.pop_front();
        if (work_unit == nullptr)
        {
            break;
        }
        
        // Resize file buffer if needed
        if (work_unit->m_file_size > file_buffer->size())
        {
            file_buffer->resize(work_unit->m_file_size);
        }
        
        // Read in the input file
        std::ifstream input_stream(work_unit->m_path.string(), 
                                   std::ios::binary | std::ios::in);
        if (! input_stream)
        {
            WRKR_OUT_ERR("unable to read from " << work_unit->m_path << "\n")
            m_errors.push_back(work_unit);
            continue;
        }
        
        input_stream.read((char*)file_buffer->data(), work_unit->m_file_size);
        input_stream.close();
        
        // Determine input, and therefore output, filetype / extension based
        // on mode
        const std::string* output_file_extension = nullptr;
        if (m_decompress)
        {
            // We are decompressing so the file we found had better be a PJG
            if ((file_buffer->at(0) == packJPG::pjg_magic[0]) && 
                (file_buffer->at(1) == packJPG::pjg_magic[1]))
            {
                output_file_extension = &JPG_EXTENSION;
            }
            else
            {
                WRKR_OUT_ERR("the input file "  << work_unit->m_path
                              << " does not appear to be a valid "
                             "packjpg (.pjg) file even though its filename "
                             "suggests it is!")
                m_errors.push_back(work_unit);
                continue;
            }
        }
        else
        {
            // We are compressing so the file we found had better be a JPG
            if ((file_buffer->at(0) == 0xFF) && (file_buffer->at(1) == 0xD8))
            {
                output_file_extension = &PJG_EXTENSION;
            }
            else
            {
                WRKR_OUT_ERR("the input file "  << work_unit->m_path
                              << " does not appear to be a valid "
                             "JPEG (.jpg) file even though its filename "
                             "suggests it is!")
                m_errors.push_back(work_unit);
                continue;
            }
        }
        
        // Do the thing!
        instance->pjglib_init_streams(file_buffer->data(), 1, 
                                      work_unit->m_file_size, nullptr, 1);
        
        bool rc = instance->pjglib_convert_stream2mem(&out_buffer, 
                                                      &out_size, message);
        if (!rc)
        {
            WRKR_OUT_ERR("An error occurred during the " 
                          <<  (m_decompress ? "decompression" : "compression")
                          << " operation on " << work_unit->m_path 
                          << ": " << message)
            m_errors.push_back(work_unit);
            continue;
        }
        
        //~ WRKR_OUT_REG("Status message: " << message)
        //~ WRKR_OUT_REG("Output size:    " << out_size)
        
        // Increment counts
        m_total_compressed_size += out_size;
        ++m_files_processed;
        
        // Create output file path be replacing extension of input file
        bfs::path out_path = bfs::path(work_unit->m_path).replace_extension(
                                *output_file_extension);
        
        // Write output file
        std::ofstream output_stream(out_path.string(), 
                                    std::ios::binary | std::ios::out);
        if (! output_stream)
        {
            WRKR_OUT_ERR("unable to write to \"" << out_path << "\"")
            m_errors.push_back(work_unit);
            continue;
        }
        output_stream.write((const char*)out_buffer, out_size);
        output_stream.close();
        
        // Delete original file if so instructed
        if (m_delete_orig)
        {
            bfs::remove(work_unit->m_path);
        }
        
        // Print status
        if (m_decompress)
        {
            WRKR_OUT_REG(" -- (" << std::fixed << std::setprecision(2) 
                         << get_global_percent_done() 
                         << "%)  Output file: " << out_path)
        }
        else
        {
            double percent = ((double)out_size) / work_unit->m_file_size;
            percent *= 100;
            WRKR_OUT_REG(" -- (" << std::fixed << std::setprecision(2) 
                         << get_global_percent_done() 
                         << "%)  Output file: " << out_path << " -- " 
                         << std::fixed << std::setprecision(2) 
                         << percent << "%")
        }
        
        delete work_unit;
    }
    
    WRKR_OUT_REG(" -- " << PURPLE << BOLD << "[[exiting]]" << ENDC)
    
    // Clean up
    std::free(out_buffer);
    delete instance;
    delete file_buffer;
}

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