#include <iostream>
#include <cstdlib>
#include <vector>
#include <string>
#include <sys/shm.h>
#include <cstring>
#include <chrono>

#include "shm_keys.hpp"
#include "row.hpp"
#include "sort_algorithms.hpp"
#include "utilization_monitor.hpp"

// Δομή Context για την πρόσβαση στον πίνακα, όπως ορίστηκε για τους schedulers
struct MatrixSortContext {
    uint8_t* A;
    int F, H, W;
};

// Δηλώσεις των συναρτήσεων των 7 Schedulers
extern long long run_static_scheduler(uint8_t* data, int F, int H, int W, int K);
extern long long run_dynamic_scheduler(MatrixSortContext* ctx, int K, int chunk_size);
extern long long run_guided_scheduler(MatrixSortContext* ctx, int K);
extern long long run_adaptive_scheduler(MatrixSortContext* ctx, int K);
extern long long run_aimd_scheduler(MatrixSortContext* ctx, int K);
extern long long run_chunk_scheduler(MatrixSortContext* ctx, int K, int chunk_size);
extern long long run_steal_scheduler(MatrixSortContext* ctx, int K, int chunk_size);

// Βοηθητική συνάρτηση για προετοιμασία των Shared Memory segments
uint8_t* attach_and_init_shm(key_t key, size_t size, uint8_t* source_data) {
    int shmid = shmget(key, size, IPC_CREAT | 0666);
    if (shmid == -1) {
        std::cerr << "Failed to get SHM for key " << key << std::endl;
        return nullptr;
    }
    uint8_t* addr = (uint8_t*)shmat(shmid, nullptr, 0);
    if (addr == (void*)-1) return nullptr;
    
    // Αντιγραφή των αρχικών δεδομένων για δίκαιο benchmark
    if (source_data) std::memcpy(addr, source_data, size);
    return addr;
}

int main(int argc, char** argv) {
    if (argc < 5) {
        std::cout << "Usage: ./shm_benchmark F H W K\n";
        return 1;
    }

    int F = std::atoi(argv[1]);
    int H = std::atoi(argv[2]);
    int W = std::atoi(argv[3]);
    int K = std::atoi(argv[4]);
    size_t totalBytes = (size_t)F * H * W;

    // 1. Σύνδεση στην αρχική μνήμη που περιέχει τον αταξινόμητο πίνακα
    int orig_shmid = shmget(SHM_KEY_ORIGINAL, totalBytes, 0666);
    if (orig_shmid == -1) {
        std::cerr << "Original SHM segment not found. Run shm_generator first!\n";
        return 1;
    }
    uint8_t* original_ptr = (uint8_t*)shmat(orig_shmid, nullptr, 0);

    // 2. Οργάνωση των 7 Schedulers για το Benchmark
    struct SchedConfig { 
        std::string name; 
        key_t key; 
        long long time; 
    };

    std::vector<SchedConfig> benchmark_list = {
        {"Static",   SHM_KEY_STATIC, 0},   {"Dynamic",  SHM_KEY_DYNAMIC, 0},
        {"Guided",   SHM_KEY_GUIDED, 0},   {"Adaptive", SHM_KEY_ADAPTIVE, 0},
        {"AIMD",     SHM_KEY_AIMD, 0},     {"Chunk",    SHM_KEY_CHUNK, 0},
        {"Steal",    SHM_KEY_STEAL, 0}
    };

    std::cout << "--- Starting Comparative Benchmark ---\n";
    std::cout << "Config: F=" << F << ", H=" << H << ", W=" << W << ", Threads=" << K << "\n\n";

    for (auto& sched : benchmark_list) {
        // Προετοιμασία μνήμης για τον συγκεκριμένο scheduler
        uint8_t* work_ptr = attach_and_init_shm(sched.key, totalBytes, original_ptr);
        MatrixSortContext ctx = {work_ptr, F, H, W};
        
        // Έναρξη παρακολούθησης πόρων
        UtilizationMonitor monitor;
        monitor.start();

        // Εκτέλεση του Scheduler
        if (sched.name == "Static") 
            sched.time = run_static_scheduler(work_ptr, F, H, W, K);
        else if (sched.name == "Dynamic")
            sched.time = run_dynamic_scheduler(&ctx, K, 10); // Chunk size 10
        else if (sched.name == "Guided")
            sched.time = run_guided_scheduler(&ctx, K);
        else if (sched.name == "Adaptive")
            sched.time = run_adaptive_scheduler(&ctx, K);
        else if (sched.name == "AIMD")
            sched.time = run_aimd_scheduler(&ctx, K);
        else if (sched.name == "Chunk")
            sched.time = run_chunk_scheduler(&ctx, K, 20); // Σταθερό chunk 20
        else if (sched.name == "Steal")
            sched.time = run_steal_scheduler(&ctx, K, 32); // Chunk 32 για κλοπή

        monitor.stop();
        
        // Εκτύπωση αποτελεσμάτων ανά αλγόριθμο
        std::cout << ">> Scheduler: " << sched.name << "\n";
        std::cout << "   Time: " << sched.time << " ms\n";
        monitor.print_stats(); 
        std::cout << "--------------------------------------\n";

        shmdt(work_ptr);
    }

    shmdt(original_ptr);
    return 0;
}