#include "parallel_tree_walker.h"
#include "tree.h"
#include <algorithm>
#include <chrono>
#include <cstring>

namespace cutysoup {

ParallelTreeWalker::ParallelTreeWalker(const Config& config) 
    : config_(config) {
}

ParallelTreeWalker::~ParallelTreeWalker() {
    shutdown_.store(true);
    work_queue_.shutdown();
    
    // Wait for any running workers to complete
    for (auto& worker : workers_) {
        if (worker.joinable()) {
            worker.join();
        }
    }
}

void ParallelTreeWalker::distribute_work(GumboNode* root, size_t estimated_nodes) {
    if (!root) return;
    
    // Clear and reset the work queue
    while (!work_queue_.empty()) {
        WorkItem dummy{nullptr, 0};
        work_queue_.try_pop(dummy);
    }
    early_termination_.store(false);
    shutdown_.store(false);
    
    // Start with the root node
    work_queue_.push(WorkItem(root, 0));
    last_stats_.work_items_created = 1;
    
    // For very large trees, pre-distribute some work
    if (estimated_nodes > config_.min_work_size * 4) {
        // Add immediate children as separate work items for better load balancing
        if (root->type == GUMBO_NODE_ELEMENT) {
            const GumboVector* children = &root->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, 1));
                last_stats_.work_items_created++;
            }
        }
    }
}

size_t ParallelTreeWalker::estimate_tree_size(GumboNode* node) const {
    if (!node) return 0;
    
    size_t count = 1;
    
    if (node->type == GUMBO_NODE_ELEMENT) {
        const GumboVector* children = &node->v.element.children;
        for (unsigned int i = 0; i < children->length; ++i) {
            // Use a simple heuristic to avoid deep recursion for estimation
            count += estimate_tree_size(static_cast<GumboNode*>(children->data[i]));
            
            // Cap estimation to avoid spending too much time on it
            if (count > 10000) {
                return count + (children->length - i - 1) * 10;  // Rough estimate for remaining
            }
        }
    }
    
    return count;
}

std::shared_ptr<Node> ParallelTreeWalker::create_node(GumboNode* gumbo_node, const Tree* tree) const {
    if (!gumbo_node || !tree) return nullptr;
    return tree->get_or_create_node(gumbo_node);
}

// Specialized optimized searches

std::vector<std::shared_ptr<Node>> ParallelTreeWalker::parallel_find_by_tag(
    GumboNode* root,
    const Tree* tree,
    const std::string& tag_name,
    const Config& config
) {
    // Convert tag name to Gumbo tag enum for faster comparison
    GumboTag target_tag = gumbo_tag_enum(tag_name.c_str());
    
    auto predicate = [target_tag, &tag_name](GumboNode* node) -> bool {
        if (!node || node->type != GUMBO_NODE_ELEMENT) {
            return false;
        }
        
        // Fast comparison using enum if possible
        if (target_tag != GUMBO_TAG_UNKNOWN) {
            return node->v.element.tag == target_tag;
        }
        
        // Fallback to string comparison for custom tags
        const char* node_name = gumbo_normalized_tagname(node->v.element.tag);
        return node_name && tag_name == node_name;
    };
    
    return parallel_find_all(root, tree, predicate, config);
}

std::vector<std::shared_ptr<Node>> ParallelTreeWalker::parallel_find_by_class(
    GumboNode* root,
    const Tree* tree,
    const std::string& class_name,
    const Config& config
) {
    auto predicate = [&class_name](GumboNode* node) -> bool {
        if (!node || node->type != GUMBO_NODE_ELEMENT) {
            return false;
        }
        
        // Look for class attribute
        const GumboVector* attributes = &node->v.element.attributes;
        for (unsigned int i = 0; i < attributes->length; ++i) {
            GumboAttribute* attr = static_cast<GumboAttribute*>(attributes->data[i]);
            if (attr && strcmp(attr->name, "class") == 0) {
                // Check if class_name is in the space-separated class list
                std::string class_attr = attr->value;
                size_t pos = 0;
                while (pos < class_attr.length()) {
                    size_t start = class_attr.find_first_not_of(" \t\r\n", pos);
                    if (start == std::string::npos) break;
                    
                    size_t end = class_attr.find_first_of(" \t\r\n", start);
                    if (end == std::string::npos) end = class_attr.length();
                    
                    std::string current_class = class_attr.substr(start, end - start);
                    if (current_class == class_name) {
                        return true;
                    }
                    
                    pos = end;
                }
                break;  // Found class attribute, no need to check others
            }
        }
        
        return false;
    };
    
    return parallel_find_all(root, tree, predicate, config);
}

std::shared_ptr<Node> ParallelTreeWalker::parallel_find_by_id(
    GumboNode* root,
    const Tree* tree,
    const std::string& id,
    const Config& config
) {
    auto predicate = [&id](GumboNode* node) -> bool {
        if (!node || node->type != GUMBO_NODE_ELEMENT) {
            return false;
        }
        
        // Look for id attribute
        const GumboVector* attributes = &node->v.element.attributes;
        for (unsigned int i = 0; i < attributes->length; ++i) {
            GumboAttribute* attr = static_cast<GumboAttribute*>(attributes->data[i]);
            if (attr && strcmp(attr->name, "id") == 0) {
                return id == attr->value;
            }
        }
        
        return false;
    };
    
    // Enable early termination for ID search since IDs should be unique
    Config id_config = config;
    id_config.early_termination = true;
    
    return parallel_find(root, tree, predicate, id_config);
}

} // namespace cutysoup
