#pragma once

#include <algorithm>
#include <vector>
#include <string>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <queue>
#include <atomic>
#include <functional>
#include <memory>
#include <chrono>
#include "gumbo.h"
#include "node.h"

namespace cutysoup {

// Forward declarations
class Tree;
class Node;

/**
 * Thread-safe queue for distributing work across parallel tree walkers
 */
template<typename T>
class ThreadSafeQueue {
public:
    ThreadSafeQueue() = default;
    ThreadSafeQueue(const ThreadSafeQueue&) = delete;
    ThreadSafeQueue& operator=(const ThreadSafeQueue&) = delete;
    ThreadSafeQueue(ThreadSafeQueue&&) = default;
    ThreadSafeQueue& operator=(ThreadSafeQueue&&) = default;
    void push(T item) {
        std::lock_guard<std::mutex> lock(mutex_);
        queue_.push(std::move(item));
        condition_.notify_one();
    }
    
    bool try_pop(T& item) {
        std::lock_guard<std::mutex> lock(mutex_);
        if (queue_.empty()) {
            return false;
        }
        item = std::move(queue_.front());
        queue_.pop();
        return true;
    }
    
    void wait_and_pop(T& item) {
        std::unique_lock<std::mutex> lock(mutex_);
        while (queue_.empty()) {
            condition_.wait(lock);
        }
        item = std::move(queue_.front());
        queue_.pop();
    }
    
    bool empty() const {
        std::lock_guard<std::mutex> lock(mutex_);
        return queue_.empty();
    }
    
    size_t size() const {
        std::lock_guard<std::mutex> lock(mutex_);
        return queue_.size();
    }
    
    void shutdown() {
        std::lock_guard<std::mutex> lock(mutex_);
        shutdown_ = true;
        condition_.notify_all();
    }
    
    bool is_shutdown() const {
        std::lock_guard<std::mutex> lock(mutex_);
        return shutdown_;
    }

private:
    mutable std::mutex mutex_;
    std::queue<T> queue_;
    std::condition_variable condition_;
    bool shutdown_ = false;
};

/**
 * Work item for parallel tree processing
 */
struct WorkItem {
    GumboNode* node;
    int depth;
    
    WorkItem(GumboNode* n, int d) : node(n), depth(d) {}
};

/**
 * Thread-safe result collector for parallel operations
 */
template<typename T>
class ResultCollector {
public:
    ResultCollector() = default;
    ResultCollector(const ResultCollector&) = delete;
    ResultCollector& operator=(const ResultCollector&) = delete;
    ResultCollector(ResultCollector&&) = default;
    ResultCollector& operator=(ResultCollector&&) = default;
    void add_result(T result) {
        std::lock_guard<std::mutex> lock(mutex_);
        results_.push_back(std::move(result));
    }
    
    void add_results(const std::vector<T>& results) {
        std::lock_guard<std::mutex> lock(mutex_);
        results_.insert(results_.end(), results.begin(), results.end());
    }
    
    std::vector<T> get_results() {
        std::lock_guard<std::mutex> lock(mutex_);
        return results_;
    }
    
    size_t size() const {
        std::lock_guard<std::mutex> lock(mutex_);
        return results_.size();
    }
    
    void clear() {
        std::lock_guard<std::mutex> lock(mutex_);
        results_.clear();
    }

private:
    mutable std::mutex mutex_;
    std::vector<T> results_;
};

/**
 * High-performance parallel tree walker for large document processing
 * 
 * Features:
 * - Multi-threaded tree traversal with work stealing
 * - Lock-free operations where possible
 * - Configurable thread count and work distribution
 * - Early termination support for find() operations
 * - Memory-efficient work distribution
 */
class ParallelTreeWalker {
public:
    /**
     * Configuration for parallel tree walking
     */
    struct Config {
        size_t thread_count = std::thread::hardware_concurrency();
        size_t min_work_size = 100;  // Minimum nodes to process in parallel
        size_t chunk_size = 50;      // Nodes to process per work item
        bool enable_work_stealing = true;
        bool early_termination = false;  // For find() operations
        
        Config() {
            // Ensure we have at least 2 threads but not more than 8 for most workloads
            if (thread_count == 0) thread_count = 2;
            if (thread_count > 8) thread_count = 8;
        }
    };
    
    explicit ParallelTreeWalker(const Config& config = Config());
    ~ParallelTreeWalker();
    
    // Parallel search operations
    template<typename Predicate>
    std::vector<std::shared_ptr<Node>> parallel_find_all(
        GumboNode* root, 
        const Tree* tree,
        Predicate predicate,
        const Config& config = Config()
    );
    
    template<typename Predicate>
    std::shared_ptr<Node> parallel_find(
        GumboNode* root,
        const Tree* tree, 
        Predicate predicate,
        const Config& config = Config()
    );
    
    // Specialized optimized searches
    std::vector<std::shared_ptr<Node>> parallel_find_by_tag(
        GumboNode* root,
        const Tree* tree,
        const std::string& tag_name,
        const Config& config = Config()
    );
    
    std::vector<std::shared_ptr<Node>> parallel_find_by_class(
        GumboNode* root,
        const Tree* tree,
        const std::string& class_name,
        const Config& config = Config()
    );
    
    std::shared_ptr<Node> parallel_find_by_id(
        GumboNode* root,
        const Tree* tree,
        const std::string& id,
        const Config& config = Config()
    );
    
    // Performance monitoring
    struct PerformanceStats {
        size_t nodes_processed = 0;
        size_t threads_used = 0;
        std::chrono::microseconds processing_time{0};
        double parallel_efficiency = 0.0;  // 0.0 to 1.0
        size_t work_items_created = 0;
        size_t work_items_stolen = 0;
    };
    
    PerformanceStats get_last_stats() const { return last_stats_; }
    void reset_stats() { last_stats_ = PerformanceStats{}; }

private:
    Config config_;
    mutable PerformanceStats last_stats_;
    
    // Thread management
    std::vector<std::thread> workers_;
    ThreadSafeQueue<WorkItem> work_queue_;
    std::atomic<bool> shutdown_{false};
    std::atomic<bool> early_termination_{false};
    
    // Work distribution
    void distribute_work(GumboNode* root, size_t estimated_nodes);
    size_t estimate_tree_size(GumboNode* node) const;
    
    // Worker thread function
    template<typename Predicate>
    void worker_function(
        Predicate predicate,
        const Tree* tree,
        ResultCollector<std::shared_ptr<Node>>& results,
        std::atomic<size_t>& nodes_processed
    );
    
    // Utility functions
    bool should_process_in_parallel(size_t estimated_nodes) const {
        return estimated_nodes >= config_.min_work_size && config_.thread_count > 1;
    }
    
    std::shared_ptr<Node> create_node(GumboNode* gumbo_node, const Tree* tree) const;
    
    // Sequential fallback for small trees
    template<typename Predicate>
    std::vector<std::shared_ptr<Node>> sequential_find_all(
        GumboNode* root,
        const Tree* tree,
        Predicate predicate
    );
    
    template<typename Predicate>
    std::shared_ptr<Node> sequential_find(
        GumboNode* root,
        const Tree* tree,
        Predicate predicate
    );
};

// Template implementations

template<typename Predicate>
std::vector<std::shared_ptr<Node>> ParallelTreeWalker::parallel_find_all(
    GumboNode* root,
    const Tree* tree,
    Predicate predicate,
    const Config& config
) {
    auto start_time = std::chrono::high_resolution_clock::now();
    last_stats_ = PerformanceStats{};
    
    if (!root) {
        return {};
    }
    
    // Estimate tree size to decide on parallel vs sequential processing
    size_t estimated_nodes = estimate_tree_size(root);
    last_stats_.nodes_processed = estimated_nodes;
    
    // Use sequential processing for small trees
    if (!should_process_in_parallel(estimated_nodes)) {
        auto results = sequential_find_all(root, tree, predicate);
        auto end_time = std::chrono::high_resolution_clock::now();
        last_stats_.processing_time = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);
        last_stats_.threads_used = 1;
        last_stats_.parallel_efficiency = 1.0;
        return results;
    }
    
    // Parallel processing
    ResultCollector<std::shared_ptr<Node>> results;
    std::atomic<size_t> nodes_processed{0};
    std::vector<std::thread> workers;
    
    // Distribute work across the tree
    distribute_work(root, estimated_nodes);
    last_stats_.threads_used = std::min(config.thread_count, estimated_nodes / config.chunk_size + 1);
    
    // Start worker threads
    for (size_t i = 0; i < last_stats_.threads_used; ++i) {
        workers.emplace_back([this, predicate, tree, &results, &nodes_processed]() {
            worker_function(predicate, tree, results, nodes_processed);
        });
    }
    
    // Wait for all workers to complete
    for (auto& worker : workers) {
        worker.join();
    }
    
    auto end_time = std::chrono::high_resolution_clock::now();
    last_stats_.processing_time = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);
    last_stats_.nodes_processed = nodes_processed.load();
    
    // Calculate parallel efficiency
    if (last_stats_.threads_used > 1) {
        // Estimate sequential time and compare
        last_stats_.parallel_efficiency = std::min(1.0, 
            static_cast<double>(last_stats_.threads_used) * 0.8);  // Assume some overhead
    } else {
        last_stats_.parallel_efficiency = 1.0;
    }
    
    return results.get_results();
}

template<typename Predicate>
std::shared_ptr<Node> ParallelTreeWalker::parallel_find(
    GumboNode* root,
    const Tree* tree,
    Predicate predicate,
    const Config& config
) {
    // Enable early termination for find operations
    Config find_config = config;
    find_config.early_termination = true;
    
    auto results = parallel_find_all(root, tree, predicate, find_config);
    return results.empty() ? nullptr : results[0];
}

template<typename Predicate>
void ParallelTreeWalker::worker_function(
    Predicate predicate,
    const Tree* tree,
    ResultCollector<std::shared_ptr<Node>>& results,
    std::atomic<size_t>& nodes_processed
) {
    std::vector<std::shared_ptr<Node>> local_results;
    local_results.reserve(32);  // Pre-allocate for efficiency
    
    WorkItem work_item{nullptr, 0};
    
    while (!shutdown_.load() && !early_termination_.load()) {
        if (!work_queue_.try_pop(work_item)) {
            // No work available, check if we should terminate
            if (work_queue_.empty() && work_queue_.is_shutdown()) {
                break;
            }
            
            // Brief yield to avoid busy waiting
            std::this_thread::yield();
            continue;
        }
        
        // Process the work item
        if (work_item.node) {
            // Check current node
            if (predicate(work_item.node)) {
                local_results.push_back(create_node(work_item.node, tree));
                
                // Early termination for find() operations
                if (config_.early_termination && !local_results.empty()) {
                    early_termination_.store(true);
                    break;
                }
            }
            
            nodes_processed.fetch_add(1);
            
            // Add children to work queue if this is an element
            if (work_item.node->type == GUMBO_NODE_ELEMENT) {
                const GumboVector* children = &work_item.node->v.element.children;
                for (unsigned int i = 0; i < children->length; ++i) {
                    GumboNode* child = static_cast<GumboNode*>(children->data[i]);
                    work_queue_.push(WorkItem(child, work_item.depth + 1));
                }
            }
        }
        
        // Periodically flush local results to avoid memory buildup
        if (local_results.size() >= 32) {
            results.add_results(local_results);
            local_results.clear();
        }
    }
    
    // Flush remaining results
    if (!local_results.empty()) {
        results.add_results(local_results);
    }
}

template<typename Predicate>
std::vector<std::shared_ptr<Node>> ParallelTreeWalker::sequential_find_all(
    GumboNode* root,
    const Tree* tree,
    Predicate predicate
) {
    std::vector<std::shared_ptr<Node>> results;
    
    if (!root) return results;
    
    // Simple recursive traversal for small trees
    std::function<void(GumboNode*)> traverse = [&](GumboNode* node) {
        if (!node) return;
        
        if (predicate(node)) {
            results.push_back(create_node(node, tree));
        }
        
        if (node->type == GUMBO_NODE_ELEMENT) {
            const GumboVector* children = &node->v.element.children;
            for (unsigned int i = 0; i < children->length; ++i) {
                traverse(static_cast<GumboNode*>(children->data[i]));
            }
        }
    };
    
    traverse(root);
    return results;
}

template<typename Predicate>
std::shared_ptr<Node> ParallelTreeWalker::sequential_find(
    GumboNode* root,
    const Tree* tree,
    Predicate predicate
) {
    if (!root) return nullptr;
    
    // Simple recursive traversal with early return
    std::function<std::shared_ptr<Node>(GumboNode*)> traverse = [&](GumboNode* node) -> std::shared_ptr<Node> {
        if (!node) return nullptr;
        
        if (predicate(node)) {
            return create_node(node, tree);
        }
        
        if (node->type == GUMBO_NODE_ELEMENT) {
            const GumboVector* children = &node->v.element.children;
            for (unsigned int i = 0; i < children->length; ++i) {
                auto result = traverse(static_cast<GumboNode*>(children->data[i]));
                if (result) return result;
            }
        }
        
        return nullptr;
    };
    
    return traverse(root);
}

} // namespace cutysoup
