#include <ldap.h>
#include <iostream>
#include <chrono>
#include <vector>
#include <string>
#include <cstring>
#include <iomanip>
#include <algorithm>
#include <thread>
#include <mutex>
#include <atomic>

using namespace std;
using namespace std::chrono;

// Configuration structure
struct LDAPConfig {
    string ldapHost;
    int ldapPort;
    string bindDN;
    string bindPassword;
    string baseDN;
    string searchUser;
    int timeout;
};

// Statistics structure for tracking operation timings (thread-safe)
struct OperationStats {
    vector<double> timings;
    string operationName;
    mutex mtx;  // Mutex for thread-safe access

    void addTiming(double ms) {
        lock_guard<mutex> lock(mtx);
        timings.push_back(ms);
    }

    void printSummary() {
        lock_guard<mutex> lock(mtx);

        if (timings.empty()) {
            cout << "\n" << operationName << ": No operations completed" << endl;
            return;
        }

        sort(timings.begin(), timings.end());

        double sum = 0;
        for (double t : timings) sum += t;
        double avg = sum / timings.size();

        double min = timings.front();
        double max = timings.back();
        double median = timings[timings.size() / 2];

        // Calculate percentiles
        double p95 = timings[(int)(timings.size() * 0.95)];
        double p99 = timings[(int)(timings.size() * 0.99)];

        cout << "\n========================================" << endl;
        cout << operationName << " Statistics:" << endl;
        cout << "========================================" << endl;
        cout << "Total Operations: " << timings.size() << endl;
        cout << "Average Time:     " << fixed << setprecision(3) << avg << " ms" << endl;
        cout << "Median Time:      " << fixed << setprecision(3) << median << " ms" << endl;
        cout << "Min Time:         " << fixed << setprecision(3) << min << " ms" << endl;
        cout << "Max Time:         " << fixed << setprecision(3) << max << " ms" << endl;
        cout << "95th Percentile:  " << fixed << setprecision(3) << p95 << " ms" << endl;
        cout << "99th Percentile:  " << fixed << setprecision(3) << p99 << " ms" << endl;
        cout << "========================================" << endl;
    }
};

class SharedLDAPBenchmark {
private:
    LDAP* ld;  // Shared LDAP connection handle
    LDAPConfig config;
    mutex ldapMutex;  // Mutex to protect LDAP operations (if needed)

public:
    SharedLDAPBenchmark(const LDAPConfig& cfg) : ld(nullptr), config(cfg) {}

    ~SharedLDAPBenchmark() {
        if (ld) {
            ldap_unbind(ld);
        }
    }

    bool initialize() {
        int rc;

        cout << "Initializing shared LDAP connection in main thread..." << endl;

        // Initialize LDAP connection (Mozilla NSLDAP style)
        ld = ldap_init((char*)config.ldapHost.c_str(), config.ldapPort);
        if (ld == nullptr) {
            cerr << "ldap_init failed" << endl;
            return false;
        }

        // Set LDAP version to 3
        int version = LDAP_VERSION3;
        rc = ldap_set_option(ld, LDAP_OPT_PROTOCOL_VERSION, &version);
        if (rc != LDAP_SUCCESS) {
            cerr << "ldap_set_option (PROTOCOL_VERSION) failed: " << ldap_err2string(rc) << endl;
            return false;
        }

        // Set timeout (in seconds)
        int timeout = config.timeout;
        ldap_set_option(ld, LDAP_OPT_TIMELIMIT, &timeout);

        // Bind to LDAP server (simple bind)
        rc = ldap_simple_bind_s(ld,
            (char*)config.bindDN.c_str(),
            (char*)config.bindPassword.c_str());
        if (rc != LDAP_SUCCESS) {
            cerr << "ldap_simple_bind_s failed: " << ldap_err2string(rc) << endl;
            return false;
        }

        cout << "Successfully connected and bound to LDAP server" << endl;
        cout << "LDAP connection handle will be shared across all threads" << endl;
        return true;
    }

    // Task 1: Simple user search
    double task1_userSearch(int threadId) {
        auto start = high_resolution_clock::now();

        string filter = "(uid=" + config.searchUser + ")";
        char* attrs[] = { (char*)"uid", nullptr };

        LDAPMessage* result = nullptr;

        // Note: LDAP operations on shared handle
        // Some LDAP implementations are thread-safe, some require external synchronization
        // Uncomment the mutex lock if you experience issues
        // lock_guard<mutex> lock(ldapMutex);

        int rc = ldap_search_ext_s(ld,
            (char*)config.baseDN.c_str(),
            LDAP_SCOPE_SUBTREE,
            (char*)filter.c_str(),
            attrs,
            0,
            nullptr,
            nullptr,
            nullptr,
            0,
            &result);

        auto end = high_resolution_clock::now();

        if (rc != LDAP_SUCCESS) {
            return -1.0;
        }

        if (result) {
            ldap_msgfree(result);
        }

        duration<double, milli> elapsed = end - start;
        return elapsed.count();
    }

    // Task 2: Fetch "uid" attribute
    double task2_fetchUidAttribute(int threadId) {
        auto start = high_resolution_clock::now();

        string filter = "uid=*";
        char* attrs[] = { (char*)"uid", nullptr };

        LDAPMessage* result = nullptr;

        // Uncomment if synchronization is needed
        // lock_guard<mutex> lock(ldapMutex);

        int rc = ldap_search_ext_s(ld,
            (char*)"cn=AAAAAA,ou=OrgUnit0,o=security.com",
            LDAP_SCOPE_SUBTREE,
            (char*)filter.c_str(),
            attrs,
            0,
            nullptr,
            nullptr,
            nullptr,
            0,
            &result);

        if (rc != LDAP_SUCCESS) {
            return -1.0;
        }

        // Count entries
        int count = ldap_count_entries(ld, result);

        if (count > 0) {
            // Get first entry
            LDAPMessage* entry = ldap_first_entry(ld, result);

            if (entry) {
                // Get first attribute
                BerElement* ber = nullptr;
                char* attr = ldap_first_attribute(ld, entry, &ber);

                if (attr) {
                    // Get values
                    char** vals = ldap_get_values(ld, entry, attr);

                    if (vals) {
                        ldap_value_free(vals);
                    }

                    ldap_memfree(attr);
                }

                if (ber) {
                    ber_free(ber, 0);
                }
            }
        }

        if (result) {
            ldap_msgfree(result);
        }

        auto end = high_resolution_clock::now();
        duration<double, milli> elapsed = end - start;
        return elapsed.count();
    }

    // Task 3: Fetch "Name" attribute (cn or displayName)
    double task3_fetchNameAttribute(int threadId) {
        auto start = high_resolution_clock::now();

        string filter = "Name=*";
        char* attrs[] = { (char*)"Name", nullptr };

        LDAPMessage* result = nullptr;

        // Uncomment if synchronization is needed
        // lock_guard<mutex> lock(ldapMutex);

        int rc = ldap_search_ext_s(ld,
            (char*)"cn=AAAAAA,ou=OrgUnit0,o=security.com",
            LDAP_SCOPE_SUBTREE,
            (char*)filter.c_str(),
            attrs,
            0,
            nullptr,
            nullptr,
            nullptr,
            0,
            &result);

        if (rc != LDAP_SUCCESS) {
            return -1.0;
        }

        // Count entries
        int count = ldap_count_entries(ld, result);

        if (count > 0) {
            // Get first entry
            LDAPMessage* entry = ldap_first_entry(ld, result);

            if (entry) {
                // Get first attribute
                BerElement* ber = nullptr;
                char* attr = ldap_first_attribute(ld, entry, &ber);

                while (attr) {
                    // Get values for this attribute
                    char** vals = ldap_get_values(ld, entry, attr);

                    if (vals) {
                        ldap_value_free(vals);
                    }

                    ldap_memfree(attr);
                    attr = ldap_next_attribute(ld, entry, ber);
                }

                if (ber) {
                    ber_free(ber, 0);
                }
            }
        }

        if (result) {
            ldap_msgfree(result);
        }

        auto end = high_resolution_clock::now();
        duration<double, milli> elapsed = end - start;
        return elapsed.count();
    }

    void runBenchmarkThread(int threadId,
        int durationMinutes,
        OperationStats& stats1,
        OperationStats& stats2,
        OperationStats& stats3,
        atomic<int>& totalIterations) {

        auto startTime = high_resolution_clock::now();
        auto endTime = startTime + minutes(durationMinutes);

        cout << "[Thread " << threadId << "] Starting benchmark using shared LDAP connection" << endl;

        int localIterations = 0;

        while (high_resolution_clock::now() < endTime) {
            localIterations++;

            // Task 1
            double time1 = task1_userSearch(threadId);
            if (time1 >= 0) {
                stats1.addTiming(time1);
            }

            // Task 2
            double time2 = task2_fetchUidAttribute(threadId);
            if (time2 >= 0) {
                stats2.addTiming(time2);
            }

            // Task 3
            double time3 = task3_fetchNameAttribute(threadId);
            if (time3 >= 0) {
                stats3.addTiming(time3);
            }
        }

        totalIterations += localIterations;
        cout << "[Thread " << threadId << "] Completed " << localIterations << " iterations" << endl;
    }
};

// Thread worker function
void threadWorker(int threadId,
    SharedLDAPBenchmark* benchmark,
    int durationMinutes,
    OperationStats& stats1,
    OperationStats& stats2,
    OperationStats& stats3,
    atomic<int>& totalIterations) {

    benchmark->runBenchmarkThread(threadId, durationMinutes, stats1, stats2, stats3, totalIterations);
}

int main(int argc, char* argv[]) {
    // Configuration - Modify these values for your LDAP server
    LDAPConfig config;
    config.ldapHost = "idpustore";  // Change to your LDAP server hostname/IP
    config.ldapPort = 22222;  // Standard LDAP port (636 for LDAPS)
    config.bindDN = "cn=admin,o=security.com";  // Change to your bind DN
    config.bindPassword = "firewall";  // Change to your password
    config.baseDN = "o=security.com";  // Change to your base DN
    config.searchUser = "AAAAAA";  // User to search for
    config.timeout = 10;  // seconds

    int numThreads = 5;
    int durationMinutes = 10;

    // Allow command-line overrides
/*    if (argc > 1) config.ldapHost = argv[1];
    if (argc > 2) config.ldapPort = atoi(argv[2]);
    if (argc > 3) config.bindDN = argv[3];
    if (argc > 4) config.bindPassword = argv[4];
    if (argc > 5) config.baseDN = argv[5];
    if (argc > 6) config.searchUser = argv[6];*/

    if (argc > 1) numThreads = atoi(argv[1]);
    if (argc > 2) durationMinutes = atoi(argv[2]);

    cout << "========================================" << endl;
    cout << "LDAP Shared Connection Performance Benchmark" << endl;
    cout << "========================================" << endl;
    cout << "LDAP Host:    " << config.ldapHost << ":" << config.ldapPort << endl;
    cout << "Base DN:      " << config.baseDN << endl;
    cout << "Search User:  " << config.searchUser << endl;
    cout << "Threads:      " << numThreads << endl;
    cout << "Duration:     " << durationMinutes << " minutes" << endl;
    cout << "Connection:   SHARED across all threads" << endl;
    cout << "========================================\n" << endl;

    // Create shared LDAP benchmark instance
    SharedLDAPBenchmark benchmark(config);

    // Initialize LDAP connection in main thread
    if (!benchmark.initialize()) {
        cerr << "Failed to initialize LDAP connection" << endl;
        return 1;
    }

    cout << "\n========================================" << endl;
    cout << "LDAP connection initialized successfully" << endl;
    cout << "All " << numThreads << " threads will share this connection" << endl;
    cout << "========================================\n" << endl;

    // Shared statistics across all threads
    OperationStats stats1, stats2, stats3;
    stats1.operationName = "Task 1: User Search (ldap_search_ext_s)";
    stats2.operationName = "Task 2: Fetch UID Attribute";
    stats3.operationName = "Task 3: Fetch Name Attribute";

    atomic<int> totalIterations(0);

    cout << "Starting " << numThreads << " worker threads..." << endl;

    auto benchmarkStart = high_resolution_clock::now();

    // Create and launch threads
    vector<thread> threads;
    for (int i = 0; i < numThreads; i++) {
        threads.emplace_back(threadWorker,
            i + 1,
            &benchmark,
            durationMinutes,
            ref(stats1),
            ref(stats2),
            ref(stats3),
            ref(totalIterations));
    }

    // Monitor progress while threads are running
    bool threadsRunning = true;
    /*thread progressMonitor([&]() {
        int lastIterations = 0;
        while (threadsRunning) {
            this_thread::sleep_for(seconds(10));
            int currentIterations = totalIterations.load();
            int iterationsPerSecond = (currentIterations - lastIterations) / 10;
            lastIterations = currentIterations;

            if (threadsRunning) {
                auto elapsed = duration_cast<seconds>(high_resolution_clock::now() - benchmarkStart);
                cout << "\n[Progress] " << elapsed.count() << "s elapsed, "
                     << currentIterations << " total iterations, "
                     << iterationsPerSecond << " ops/sec (last 10s)" << endl;
            }
        }
    });*/

    // Wait for all threads to complete
    for (auto& t : threads) {
        t.join();
    }

    threadsRunning = false;
    //progressMonitor.join();

    auto benchmarkEnd = high_resolution_clock::now();
    auto totalDuration = duration_cast<seconds>(benchmarkEnd - benchmarkStart);

    cout << "\n========================================" << endl;
    cout << "Benchmark Completed!" << endl;
    cout << "========================================" << endl;
    cout << "Threads:              " << numThreads << endl;
    cout << "Connection Type:      SHARED" << endl;
    cout << "Total Duration:       " << totalDuration.count() << " seconds" << endl;
    cout << "Total Iterations:     " << totalIterations.load() << endl;
    cout << "Overall Throughput:   " << fixed << setprecision(2)
        << (totalIterations.load() * 3.0 / totalDuration.count()) << " ops/sec" << endl;
    cout << "========================================" << endl;

    // Print summaries
    stats1.printSummary();
    stats2.printSummary();
    stats3.printSummary();

    return 0;
}