#pragma once

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

namespace cutysoup {

// Forward declarations
class Tree;
class Node;

/**
 * Simple, robust parallel tree walker that avoids complex coordination
 * 
 * Design principles:
 * - Pre-collect all nodes to avoid dynamic work distribution
 * - Use simple thread coordination without complex queues
 * - Clear timeout and fallback mechanisms
 * - Minimal shared state to reduce race conditions
 */
class SimpleParallelWalker {
public:
    struct Config {
        size_t thread_count = std::thread::hardware_concurrency();
        size_t min_nodes_for_parallel = 200;  // Minimum nodes to justify parallel processing
        std::chrono::milliseconds timeout{5000};  // 5 second timeout
        
        Config() {
            // Ensure reasonable thread count
            if (thread_count == 0) thread_count = 2;
            if (thread_count > 8) thread_count = 8;  // Avoid over-threading for most workloads
        }
    };
    
    explicit SimpleParallelWalker(const Config& config = Config());
    ~SimpleParallelWalker() = default;
    
    // Main 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 fast searches
    std::vector<std::shared_ptr<Node>> find_by_tag_parallel(
        GumboNode* root,
        const Tree* tree,
        const std::string& tag_name,
        const Config& config = Config()
    );
    
    std::vector<std::shared_ptr<Node>> find_by_class_parallel(
        GumboNode* root,
        const Tree* tree,
        const std::string& class_name,
        const Config& config = Config()
    );
    
    std::shared_ptr<Node> find_by_id_parallel(
        GumboNode* root,
        const Tree* tree,
        const std::string& id,
        const Config& config = Config()
    );

private:
    Config config_;
    
    // Pre-collect all nodes in the tree
    std::vector<GumboNode*> collect_all_nodes(GumboNode* root) const;
    
    // Distribute nodes across threads for processing
    std::vector<std::vector<GumboNode*>> distribute_nodes(
        const std::vector<GumboNode*>& all_nodes,
        size_t thread_count
    ) const;
    
    // Worker function for parallel processing
    template<typename Predicate>
    void process_node_chunk(
        const std::vector<GumboNode*>& nodes,
        const Tree* tree,
        Predicate predicate,
        std::vector<std::shared_ptr<Node>>& results,
        std::mutex& results_mutex,
        std::atomic<bool>& early_termination
    ) const;
    
    // Helper to create Node from GumboNode
    std::shared_ptr<Node> create_node(GumboNode* gumbo_node, const Tree* tree) const;
    
    // Sequential fallback
    template<typename Predicate>
    std::vector<std::shared_ptr<Node>> sequential_find_all(
        GumboNode* root,
        const Tree* tree,
        Predicate predicate
    ) const;
    
    template<typename Predicate>
    std::shared_ptr<Node> sequential_find(
        GumboNode* root,
        const Tree* tree,
        Predicate predicate
    ) const;
};

// Template implementations

template<typename Predicate>
std::vector<std::shared_ptr<Node>> SimpleParallelWalker::parallel_find_all(
    GumboNode* root,
    const Tree* tree,
    Predicate predicate,
    const Config& config
) {
    if (!root || !tree) {
        return {};
    }
    
    // Collect all nodes first
    auto all_nodes = collect_all_nodes(root);
    
    // Use sequential processing for small trees
    if (all_nodes.size() < config.min_nodes_for_parallel || config.thread_count <= 1) {
        return sequential_find_all(root, tree, predicate);
    }
    
    // Distribute nodes across threads
    auto node_chunks = distribute_nodes(all_nodes, config.thread_count);
    
    // Prepare shared state
    std::vector<std::shared_ptr<Node>> final_results;
    std::mutex results_mutex;
    std::atomic<bool> early_termination{false};
    
    // Start worker threads
    std::vector<std::thread> workers;
    std::vector<std::vector<std::shared_ptr<Node>>> thread_results(node_chunks.size());
    
    for (size_t i = 0; i < node_chunks.size(); ++i) {
        workers.emplace_back([this, &node_chunks, i, tree, predicate, &thread_results, &results_mutex, &early_termination]() {
            process_node_chunk(node_chunks[i], tree, predicate, thread_results[i], results_mutex, early_termination);
        });
    }
    
    // Wait for all workers with timeout
    auto start_time = std::chrono::steady_clock::now();
    for (auto& worker : workers) {
        if (worker.joinable()) {
            // Simple join - if it hangs, we'll timeout at the caller level
            worker.join();
        }
        
        // Check timeout
        auto elapsed = std::chrono::steady_clock::now() - start_time;
        if (elapsed > config.timeout) {
            // Timeout - terminate early and fall back to sequential
            early_termination.store(true);
            
            // Detach remaining threads and fall back to sequential
            for (auto& remaining_worker : workers) {
                if (remaining_worker.joinable()) {
                    remaining_worker.detach();
                }
            }
            
            return sequential_find_all(root, tree, predicate);
        }
    }
    
    // Combine results from all threads
    for (const auto& thread_result : thread_results) {
        final_results.insert(final_results.end(), thread_result.begin(), thread_result.end());
    }
    
    return final_results;
}

template<typename Predicate>
std::shared_ptr<Node> SimpleParallelWalker::parallel_find(
    GumboNode* root,
    const Tree* tree,
    Predicate predicate,
    const Config& config
) {
    if (!root || !tree) {
        return nullptr;
    }
    
    // For find (single result), use a modified approach with early termination
    auto all_nodes = collect_all_nodes(root);
    
    if (all_nodes.size() < config.min_nodes_for_parallel || config.thread_count <= 1) {
        return sequential_find(root, tree, predicate);
    }
    
    auto node_chunks = distribute_nodes(all_nodes, config.thread_count);
    
    std::atomic<bool> found{false};
    std::shared_ptr<Node> result{nullptr};
    std::mutex result_mutex;
    
    std::vector<std::thread> workers;
    
    for (size_t i = 0; i < node_chunks.size(); ++i) {
        workers.emplace_back([this, &node_chunks, i, tree, predicate, &found, &result, &result_mutex]() {
            for (GumboNode* node : node_chunks[i]) {
                if (found.load()) {
                    break;  // Another thread found the result
                }
                
                if (predicate(node)) {
                    std::lock_guard<std::mutex> lock(result_mutex);
                    if (!found.load()) {  // Double-check under lock
                        result = create_node(node, tree);
                        found.store(true);
                        break;
                    }
                }
            }
        });
    }
    
    // Wait for workers with timeout
    auto start_time = std::chrono::steady_clock::now();
    for (auto& worker : workers) {
        if (worker.joinable()) {
            worker.join();
        }
        
        auto elapsed = std::chrono::steady_clock::now() - start_time;
        if (elapsed > config.timeout) {
            found.store(true);  // Signal other threads to stop
            
            for (auto& remaining_worker : workers) {
                if (remaining_worker.joinable()) {
                    remaining_worker.detach();
                }
            }
            
            return sequential_find(root, tree, predicate);
        }
    }
    
    return result;
}

template<typename Predicate>
void SimpleParallelWalker::process_node_chunk(
    const std::vector<GumboNode*>& nodes,
    const Tree* tree,
    Predicate predicate,
    std::vector<std::shared_ptr<Node>>& results,
    std::mutex& results_mutex,
    std::atomic<bool>& early_termination
) const {
    std::vector<std::shared_ptr<Node>> local_results;
    
    for (GumboNode* node : nodes) {
        if (early_termination.load()) {
            break;
        }
        
        if (predicate(node)) {
            local_results.push_back(create_node(node, tree));
        }
    }
    
    // Update shared results
    if (!local_results.empty() && !early_termination.load()) {
        std::lock_guard<std::mutex> lock(results_mutex);
        results.insert(results.end(), local_results.begin(), local_results.end());
    }
}

template<typename Predicate>
std::vector<std::shared_ptr<Node>> SimpleParallelWalker::sequential_find_all(
    GumboNode* root,
    const Tree* tree,
    Predicate predicate
) const {
    std::vector<std::shared_ptr<Node>> results;
    
    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> SimpleParallelWalker::sequential_find(
    GumboNode* root,
    const Tree* tree,
    Predicate predicate
) const {
    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
