#include "utilization_monitor.hpp"

#include <algorithm>
#include <fstream>
#include <sstream>
#include <string>

UtilizationMonitor::UtilizationMonitor(unsigned int interval_seconds,
                                       double a,
                                       double beta)
    : stop_flag(false),
      running(false),
      interval_seconds(interval_seconds),
      previous(0),
      current(0),
      cpu_count(0),
      per_cpu_mean_util(0),
      per_cpu_mean_count(0),
      a(a),
      beta(beta),
      total_util_ewma(0.0),
      total_util_ewma_raw(0.0),
      ewma_initialized(false) {
}

UtilizationMonitor::~UtilizationMonitor() {
    stop();
    freeMatrix(previous, cpu_count);
    freeMatrix(current, cpu_count);
    freeDoubleArray(per_cpu_mean_util);
}

void UtilizationMonitor::start() {
    if (running.load()) {
        return;
    }

    stop_flag.store(false);
    running.store(true);
    worker_thread = boost::thread(&UtilizationMonitor::monitorLoop, this);
}

void UtilizationMonitor::stop() {
    stop_flag.store(true);

    if (worker_thread.joinable()) {
        worker_thread.join();
    }

    running.store(false);
}

bool UtilizationMonitor::isRunning() {
    return running.load();
}

double UtilizationMonitor::getEwmaUtilization() {
    boost::lock_guard<boost::mutex> lock(data_mutex);
    return total_util_ewma;
}

double UtilizationMonitor::getEwmaRawUtilization() {
    boost::lock_guard<boost::mutex> lock(data_mutex);
    return total_util_ewma_raw;
}
unsigned long long** UtilizationMonitor::allocateMatrix(int rows) {
    unsigned long long** matrix = new unsigned long long*[rows];

    for (int i = 0; i < rows; ++i) {
        matrix[i] = new unsigned long long[COL_COUNT];
        for (int j = 0; j < COL_COUNT; ++j) {
            matrix[i][j] = 0;
        }
    }

    return matrix;
}

void UtilizationMonitor::freeMatrix(unsigned long long** matrix, int rows) {
    if (!matrix) {
        return;
    }

    for (int i = 0; i < rows; ++i) {
        delete[] matrix[i];
    }

    delete[] matrix;
}

double* UtilizationMonitor::allocateDoubleArray(int n) {
    double* arr = new double[n];

    for (int i = 0; i < n; ++i) {
        arr[i] = 0.0;
    }

    return arr;
}

void UtilizationMonitor::freeDoubleArray(double* arr) {
    if (arr) {
        delete[] arr;
    }
}

void UtilizationMonitor::ensureMeanStorage(int n) {
    if (per_cpu_mean_util && per_cpu_mean_count == n) {
        return;
    }

    freeDoubleArray(per_cpu_mean_util);
    per_cpu_mean_util = allocateDoubleArray(n);
    per_cpu_mean_count = n;
}

unsigned long long UtilizationMonitor::getRunningTime(unsigned long long* row) {
    return row[COL_USER] + row[COL_NICE];
}

unsigned long long UtilizationMonitor::getIoWaitTime(unsigned long long* row) {
    return row[COL_IOWAIT];
}

unsigned long long UtilizationMonitor::getIdleTime(unsigned long long* row) {
    return row[COL_IDLE];
}

unsigned long long UtilizationMonitor::getTotalTime(unsigned long long* row) {
    return row[COL_USER]
         + row[COL_NICE]
         + row[COL_SYSTEM]
         + row[COL_IDLE]
         + row[COL_IOWAIT]
         + row[COL_IRQ]
         + row[COL_SOFTIRQ]
         + row[COL_STEAL]
         + row[COL_GUEST]
         + row[COL_GUEST_NICE];
}

unsigned long long UtilizationMonitor::getBusyTime(unsigned long long* row) {
    return row[COL_USER]
         + row[COL_NICE]
         + row[COL_SYSTEM]
         + row[COL_IOWAIT]
         + row[COL_IRQ]
         + row[COL_SOFTIRQ]
         + row[COL_STEAL]
         + row[COL_GUEST]
         + row[COL_GUEST_NICE];
}

bool UtilizationMonitor::readPerCpuStats(unsigned long long**& matrix,
                                         int& cpu_count) {
    std::ifstream file("/proc/stat");
    if (!file.is_open()) {
        return false;
    }

    std::string line;
    int detected_cpus = 0;

    while (std::getline(file, line)) {
        if (line.rfind("cpu", 0) != 0) {
            continue;
        }

        if (line.size() > 3 && line[3] == ' ') {
            continue;
        }

        ++detected_cpus;
    }

    if (detected_cpus == 0) {
        return false;
    }

    file.clear();
    file.seekg(0);

    unsigned long long** new_matrix = allocateMatrix(detected_cpus);
    int row = 0;

    while (std::getline(file, line)) {
        if (line.rfind("cpu", 0) != 0) {
            continue;
        }

        if (line.size() > 3 && line[3] == ' ') {
            continue;
        }

        std::istringstream iss(line);
        std::string cpu_label;

        iss >> cpu_label
            >> new_matrix[row][COL_USER]
            >> new_matrix[row][COL_NICE]
            >> new_matrix[row][COL_SYSTEM]
            >> new_matrix[row][COL_IDLE]
            >> new_matrix[row][COL_IOWAIT]
            >> new_matrix[row][COL_IRQ]
            >> new_matrix[row][COL_SOFTIRQ]
            >> new_matrix[row][COL_STEAL]
            >> new_matrix[row][COL_GUEST]
            >> new_matrix[row][COL_GUEST_NICE];

        if (iss.fail()) {
            freeMatrix(new_matrix, detected_cpus);
            return false;
        }

        ++row;
        if (row >= detected_cpus) {
            break;
        }
    }

    freeMatrix(matrix, cpu_count);
    matrix = new_matrix;
    cpu_count = detected_cpus;

    return true;
}

UtilizationMonitor::CpuUsageResult
UtilizationMonitor::calculateUsage(unsigned long long* prev_row,
                                   unsigned long long* curr_row,
                                   int cpu_index) {
    CpuUsageResult r;

    r.cpu_index = cpu_index;
    r.running_delta = getRunningTime(curr_row) - getRunningTime(prev_row);
    r.iowait_delta = getIoWaitTime(curr_row) - getIoWaitTime(prev_row);
    r.idle_delta = getIdleTime(curr_row) - getIdleTime(prev_row);
    r.total_delta = getTotalTime(curr_row) - getTotalTime(prev_row);

    unsigned long long busy_delta = getBusyTime(curr_row) - getBusyTime(prev_row);

    r.running_percent = 0.0;
    r.iowait_percent = 0.0;
    r.idle_percent = 0.0;
    r.utilization_percent = 0.0;
    r.mean_utilization_percent = 0.0;

    if (r.total_delta > 0) {
        r.running_percent = 100.0 * (double) r.running_delta / (double) r.total_delta;
        r.iowait_percent = 100.0 * (double) r.iowait_delta / (double) r.total_delta;
        r.idle_percent = 100.0 * (double) r.idle_delta / (double) r.total_delta;
        r.utilization_percent = 100.0 * (double) busy_delta / (double) r.total_delta;
    }

    return r;
}

void UtilizationMonitor::updatePerCpuMeans(CpuUsageResult* results,
                                           int count,
                                           unsigned long long sample_index) {
    for (int i = 0; i < count; ++i) {
        if (sample_index == 1) {
            per_cpu_mean_util[i] = results[i].utilization_percent;
        } else {
            per_cpu_mean_util[i] =
                ((per_cpu_mean_util[i] * (double) (sample_index - 1))
                 + results[i].utilization_percent)
                / (double) sample_index;
        }

        results[i].mean_utilization_percent = per_cpu_mean_util[i];
    }
}

double UtilizationMonitor::computeTotalUtilizationEwma(CpuUsageResult* results,
                                                       int count, bool raw=false) {
    if (count == 0) {
        return 0.0;
    }

    double max_mean = results[0].mean_utilization_percent;
    double min_current = results[0].utilization_percent;

    for (int i = 1; i < count; ++i) {
        if (results[i].mean_utilization_percent > max_mean) {
            max_mean = results[i].mean_utilization_percent;
        }

        if (results[i].utilization_percent < min_current) {
            min_current = results[i].utilization_percent;
        }
    }

    double combined = a * max_mean + (1.0 - a) * min_current;

    boost::lock_guard<boost::mutex> lock(data_mutex);

    if (!ewma_initialized) {
        total_util_ewma = combined;
	total_util_ewma_raw=combined;
        ewma_initialized = true;
    } else {
	total_util_ewma_raw=combined;
        total_util_ewma = beta * total_util_ewma + (1.0 - beta) * total_util_ewma_raw;
    }
    if (raw) 
	return total_util_ewma_raw;
    else
    	return total_util_ewma;
}

void UtilizationMonitor::monitorLoop() {
    if (!readPerCpuStats(previous, cpu_count)) {
        running.store(false);
        return;
    }

    ensureMeanStorage(cpu_count);

    unsigned long long sample_index = 0;

    while (!stop_flag.load()) {
        for (unsigned int i = 0; i < interval_seconds; ++i) {
            if (stop_flag.load()) {
                running.store(false);
                return;
            }

            boost::this_thread::sleep_for(boost::chrono::seconds(1));
        }

        int current_cpu_count = 0;
        if (!readPerCpuStats(current, current_cpu_count)) {
            continue;
        }

        int n = std::min(cpu_count, current_cpu_count);
        ensureMeanStorage(n);

        CpuUsageResult* results = new CpuUsageResult[n];

        for (int cpu = 0; cpu < n; ++cpu) {
            results[cpu] = calculateUsage(previous[cpu], current[cpu], cpu);
        }

        ++sample_index;
        updatePerCpuMeans(results, n, sample_index);
        computeTotalUtilizationEwma(results, n);

        delete[] results;

        freeMatrix(previous, cpu_count);
        previous = current;
        cpu_count = current_cpu_count;
        current = 0;
    }

    running.store(false);
}