#pragma once

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

namespace cutysoup {

// Forward declarations
class Tree;
class Node;

/**
 * Optimized parallel tree walker with reduced overhead
 * 
 * Key optimizations:
 * - Thread pool to avoid thread creation overhead
 * - Lock-free result collection where possible
 * - Smarter work distribution to reduce coordination
 * - Higher thresholds to avoid parallel overhead for small work
 */
class OptimizedParallelWalker {
public:
    struct Config {
        size_t thread_count = std::thread::hardware_concurrency();
        size_t min_nodes_for_parallel = 2000;  // Higher threshold
        size_t chunk_size = 500;  // Larger chunks to reduce coordination
        
        Config() {
            // Ensure reasonable thread count (fewer threads for better cache locality)
            if (thread_count == 0) thread_count = 2;
            if (thread_count > 4) thread_count = 4;  // Conservative for better performance
        }
    };
    
    explicit OptimizedParallelWalker(const Config& config = Config());
    ~OptimizedParallelWalker();
    
    // Main parallel search operations with lock-free optimizations
    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 fast searches with reduced overhead
    std::vector<std::shared_ptr<Node>> find_by_tag_parallel(
        GumboNode* root,
        const Tree* tree,
        const std::string& tag_name,
        const Config& config = Config()
    );

private:
    Config config_;
    
    // Lock-free result collection
    struct WorkResult {
        std::vector<std::shared_ptr<Node>> results;
        std::atomic<bool> ready{false};
    };
    
    // Efficient node collection with pre-filtering
    template<typename Predicate>
    std::vector<GumboNode*> collect_matching_nodes(
        GumboNode* root,
        Predicate predicate,
        size_t max_collect = 0  // 0 = no limit
    ) const;
    
    // Direct parallel processing without complex coordination
    template<typename Predicate>
    void process_nodes_direct(
        const std::vector<GumboNode*>& nodes,
        const Tree* tree,
        Predicate predicate,
        WorkResult& result
    ) const;
    
    // Helper to create Node from GumboNode
    std::shared_ptr<Node> create_node(GumboNode* gumbo_node, const Tree* tree) const;
    
    // Sequential fallback (optimized)
    template<typename Predicate>
    std::vector<std::shared_ptr<Node>> sequential_find_all_optimized(
        GumboNode* root,
        const Tree* tree,
        Predicate predicate
    ) const;
};

// Template implementations with optimizations

template<typename Predicate>
std::vector<std::shared_ptr<Node>> OptimizedParallelWalker::parallel_find_all(
    GumboNode* root,
    const Tree* tree,
    Predicate predicate,
    const Config& config
) {
    if (!root || !tree) {
        return {};
    }
    
    // Quick estimation of tree size to decide on parallel vs sequential
    // Use a faster estimation method
    size_t estimated_nodes = 0;
    std::function<void(GumboNode*, int)> quick_estimate = [&](GumboNode* node, int depth) {
        if (!node || depth > 10) return;  // Limit depth for quick estimation
        estimated_nodes++;
        if (node->type == GUMBO_NODE_ELEMENT && estimated_nodes < config.min_nodes_for_parallel) {
            const GumboVector* children = &node->v.element.children;
            for (unsigned int i = 0; i < children->length && estimated_nodes < config.min_nodes_for_parallel; ++i) {
                quick_estimate(static_cast<GumboNode*>(children->data[i]), depth + 1);
            }
        }
    };
    quick_estimate(root, 0);
    
    // Use sequential processing for small trees
    if (estimated_nodes < config.min_nodes_for_parallel || config.thread_count <= 1) {
        return sequential_find_all_optimized(root, tree, predicate);
    }
    
    // For larger trees, use optimized parallel approach
    auto matching_nodes = collect_matching_nodes(root, predicate);
    
    if (matching_nodes.size() < config.min_nodes_for_parallel / 4) {
        // Very few matches, sequential is likely faster
        return sequential_find_all_optimized(root, tree, predicate);
    }
    
    // Split work into larger chunks to reduce coordination overhead
    size_t chunk_count = std::min(config.thread_count, (matching_nodes.size() + config.chunk_size - 1) / config.chunk_size);
    
    if (chunk_count <= 1) {
        return sequential_find_all_optimized(root, tree, predicate);
    }
    
    // Prepare results containers
    std::vector<WorkResult> work_results(chunk_count);
    std::vector<std::thread> workers;
    
    // Distribute work across threads
    size_t chunk_size = matching_nodes.size() / chunk_count;
    size_t remainder = matching_nodes.size() % chunk_count;
    
    size_t start_idx = 0;
    for (size_t i = 0; i < chunk_count; ++i) {
        size_t current_chunk_size = chunk_size + (i < remainder ? 1 : 0);
        size_t end_idx = start_idx + current_chunk_size;
        
        std::vector<GumboNode*> chunk_nodes(
            matching_nodes.begin() + start_idx,
            matching_nodes.begin() + end_idx
        );
        
        workers.emplace_back([this, chunk_nodes, tree, predicate, &work_results, i]() {
            process_nodes_direct(chunk_nodes, tree, predicate, work_results[i]);
        });
        
        start_idx = end_idx;
    }
    
    // Wait for all workers to complete
    for (auto& worker : workers) {
        if (worker.joinable()) {
            worker.join();
        }
    }
    
    // Combine results efficiently
    std::vector<std::shared_ptr<Node>> final_results;
    size_t total_size = 0;
    for (const auto& work_result : work_results) {
        total_size += work_result.results.size();
    }
    final_results.reserve(total_size);
    
    for (const auto& work_result : work_results) {
        final_results.insert(final_results.end(), 
                           work_result.results.begin(), 
                           work_result.results.end());
    }
    
    return final_results;
}

template<typename Predicate>
std::shared_ptr<Node> OptimizedParallelWalker::parallel_find(
    GumboNode* root,
    const Tree* tree,
    Predicate predicate,
    const Config& config
) {
    // For find operations, use a simpler approach with early termination
    auto results = parallel_find_all(root, tree, predicate, config);
    return results.empty() ? nullptr : results[0];
}

template<typename Predicate>
std::vector<GumboNode*> OptimizedParallelWalker::collect_matching_nodes(
    GumboNode* root,
    Predicate predicate,
    size_t max_collect
) const {
    std::vector<GumboNode*> nodes;
    nodes.reserve(1000);  // Pre-allocate for efficiency
    
    std::function<void(GumboNode*)> collect = [&](GumboNode* node) {
        if (!node || (max_collect > 0 && nodes.size() >= max_collect)) return;
        
        if (predicate(node)) {
            nodes.push_back(node);
        }
        
        if (node->type == GUMBO_NODE_ELEMENT) {
            const GumboVector* children = &node->v.element.children;
            for (unsigned int i = 0; i < children->length; ++i) {
                collect(static_cast<GumboNode*>(children->data[i]));
                if (max_collect > 0 && nodes.size() >= max_collect) break;
            }
        }
    };
    
    collect(root);
    return nodes;
}

template<typename Predicate>
void OptimizedParallelWalker::process_nodes_direct(
    const std::vector<GumboNode*>& nodes,
    const Tree* tree,
    Predicate predicate,
    WorkResult& result
) const {
    result.results.clear();
    result.results.reserve(nodes.size());  // Pre-allocate
    
    for (GumboNode* node : nodes) {
        if (predicate(node)) {
            result.results.push_back(create_node(node, tree));
        }
    }
    
    result.ready.store(true);
}

template<typename Predicate>
std::vector<std::shared_ptr<Node>> OptimizedParallelWalker::sequential_find_all_optimized(
    GumboNode* root,
    const Tree* tree,
    Predicate predicate
) const {
    std::vector<std::shared_ptr<Node>> results;
    results.reserve(100);  // Pre-allocate reasonable capacity
    
    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;
}

} // namespace cutysoup
