#include "ultra_fast_walker.h"
#include "tree.h"
#include <algorithm>
#include <cstring>

namespace cutysoup {

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

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

const char* UltraFastWalker::find_attribute_fast(const GumboVector* attributes, const char* attr_name) const {
    if (unlikely(!attributes || !attr_name)) return nullptr;
    
    const size_t attr_name_len = strlen(attr_name);
    
    // Manually unroll loop for common cases
    for (unsigned int i = 0; likely(i < attributes->length); ++i) {
        GumboAttribute* attr = static_cast<GumboAttribute*>(attributes->data[i]);
        if (likely(attr && attr->name)) {
            // Fast string comparison using SIMD if enabled and appropriate
            if (config_.enable_simd && attr_name_len >= 8) {
                if (simd_string_equals(attr->name, attr_name, attr_name_len)) {
                    return attr->value;
                }
            } else {
                // Fast memcmp for short strings
                if (strlen(attr->name) == attr_name_len && 
                    memcmp(attr->name, attr_name, attr_name_len) == 0) {
                    return attr->value;
                }
            }
        }
    }
    
    return nullptr;
}

bool UltraFastWalker::has_class_fast(const GumboVector* attributes, const char* class_name, size_t class_len) const {
    const char* class_attr = find_attribute_fast(attributes, "class");
    if (unlikely(!class_attr)) return false;
    
    const size_t attr_len = strlen(class_attr);
    
    if (config_.enable_simd && class_len >= 4 && attr_len >= 8) {
        return simd_contains_class(class_attr, class_name, attr_len, class_len);
    }
    
    // Optimized sequential search for class names
    const char* pos = class_attr;
    const char* end = class_attr + attr_len;
    
    while (likely(pos < end)) {
        // Skip whitespace
        while (likely(pos < end && (*pos == ' ' || *pos == '\t' || *pos == '\r' || *pos == '\n'))) {
            ++pos;
        }
        
        if (unlikely(pos >= end)) break;
        
        // Find end of current class
        const char* class_start = pos;
        while (likely(pos < end && *pos != ' ' && *pos != '\t' && *pos != '\r' && *pos != '\n')) {
            ++pos;
        }
        
        const size_t current_class_len = pos - class_start;
        
        // Quick length check before comparison
        if (likely(current_class_len == class_len)) {
            if (memcmp(class_start, class_name, class_len) == 0) {
                return true;
            }
        }
    }
    
    return false;
}

bool UltraFastWalker::simd_string_equals(const char* str1, const char* str2, size_t len) const {
    if (unlikely(!str1 || !str2)) return false;
    
#ifdef __AVX2__
    // AVX2 optimized comparison for strings >= 32 bytes
    if (len >= 32) {
        const size_t simd_chunks = len / 32;
        for (size_t i = 0; i < simd_chunks; ++i) {
            __m256i v1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(str1 + i * 32));
            __m256i v2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(str2 + i * 32));
            __m256i cmp = _mm256_cmpeq_epi8(v1, v2);
            
            if (_mm256_movemask_epi8(cmp) != 0xFFFFFFFF) {
                return false;
            }
        }
        
        // Handle remaining bytes
        const size_t remaining = len % 32;
        if (remaining > 0) {
            return memcmp(str1 + simd_chunks * 32, str2 + simd_chunks * 32, remaining) == 0;
        }
        
        return true;
    }
#endif

#ifdef __SSE2__
    // SSE2 optimized comparison for strings >= 16 bytes
    if (len >= 16) {
        const size_t simd_chunks = len / 16;
        for (size_t i = 0; i < simd_chunks; ++i) {
            __m128i v1 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(str1 + i * 16));
            __m128i v2 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(str2 + i * 16));
            __m128i cmp = _mm_cmpeq_epi8(v1, v2);
            
            if (_mm_movemask_epi8(cmp) != 0xFFFF) {
                return false;
            }
        }
        
        // Handle remaining bytes
        const size_t remaining = len % 16;
        if (remaining > 0) {
            return memcmp(str1 + simd_chunks * 16, str2 + simd_chunks * 16, remaining) == 0;
        }
        
        return true;
    }
#endif
    
    // Fallback to memcmp for short strings
    return memcmp(str1, str2, len) == 0;
}

bool UltraFastWalker::simd_contains_class(const char* class_attr, const char* target_class, 
                                        size_t attr_len, size_t target_len) const {
    if (unlikely(!class_attr || !target_class || target_len == 0)) return false;
    
#ifdef __AVX2__
    if (target_len <= 32 && attr_len >= 32) {
        // Load target class into SIMD register
        char target_buffer[32] = {0};
        memcpy(target_buffer, target_class, target_len);
        __m256i target_vec = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(target_buffer));
        
        // Scan through the class attribute
        for (size_t i = 0; i <= attr_len - 32; i += 32) {
            __m256i attr_vec = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(class_attr + i));
            __m256i cmp = _mm256_cmpeq_epi8(attr_vec, target_vec);
            
            uint32_t mask = _mm256_movemask_epi8(cmp);
            if (mask != 0) {
                // Potential match found, verify boundaries
                for (int bit = 0; bit < 32; ++bit) {
                    if (mask & (1u << bit)) {
                        size_t pos = i + bit;
                        if (pos + target_len <= attr_len) {
                            // Check if this is a word boundary match
                            bool start_ok = (pos == 0 || class_attr[pos - 1] == ' ' || 
                                           class_attr[pos - 1] == '\t' || class_attr[pos - 1] == '\n');
                            bool end_ok = (pos + target_len == attr_len || 
                                         class_attr[pos + target_len] == ' ' ||
                                         class_attr[pos + target_len] == '\t' ||
                                         class_attr[pos + target_len] == '\n');
                            
                            if (start_ok && end_ok && 
                                memcmp(class_attr + pos, target_class, target_len) == 0) {
                                return true;
                            }
                        }
                    }
                }
            }
        }
    }
#endif
    
    // Fallback to standard search
    return has_class_fast(nullptr, target_class, target_len);  // This will use the non-SIMD path
}

std::vector<std::shared_ptr<Node>> UltraFastWalker::find_by_tag_ultra_fast(
    GumboNode* root,
    const Tree* tree,
    const std::string& tag_name
) {
    std::vector<std::shared_ptr<Node>> results;
    results.reserve(config_.reserve_capacity);
    
    if (unlikely(!root)) return results;
    
    // Convert tag name to Gumbo tag enum for ultra-fast comparison
    GumboTag target_tag = gumbo_tag_enum(tag_name.c_str());
    
    if (likely(target_tag != GUMBO_TAG_UNKNOWN)) {
        // Optimized path for known tags - use enum comparison
        traverse_depth_first_optimized(root, [&](GumboNode* node) {
            if (likely(node && node->type == GUMBO_NODE_ELEMENT)) {
                if (likely(node->v.element.tag == target_tag)) {
                    results.emplace_back(create_node_fast(node, tree));
                }
            }
        });
    } else {
        // Fallback for custom tags - use SIMD-optimized string comparison
        const char* tag_cstr = tag_name.c_str();
        const size_t tag_len = tag_name.length();
        
        traverse_depth_first_optimized(root, [&](GumboNode* node) {
            if (likely(node && node->type == GUMBO_NODE_ELEMENT)) {
                const char* node_name = gumbo_normalized_tagname(node->v.element.tag);
                if (likely(node_name)) {
                    if (config_.enable_simd && tag_len >= 8) {
                        if (simd_string_equals(node_name, tag_cstr, tag_len)) {
                            results.emplace_back(create_node_fast(node, tree));
                        }
                    } else {
                        if (strlen(node_name) == tag_len && memcmp(node_name, tag_cstr, tag_len) == 0) {
                            results.emplace_back(create_node_fast(node, tree));
                        }
                    }
                }
            }
        });
    }
    
    return results;
}

std::vector<std::shared_ptr<Node>> UltraFastWalker::find_by_class_ultra_fast(
    GumboNode* root,
    const Tree* tree,
    const std::string& class_name
) {
    std::vector<std::shared_ptr<Node>> results;
    results.reserve(config_.reserve_capacity);
    
    if (unlikely(!root)) return results;
    
    const char* class_cstr = class_name.c_str();
    const size_t class_len = class_name.length();
    
    traverse_depth_first_optimized(root, [&](GumboNode* node) {
        if (likely(node && node->type == GUMBO_NODE_ELEMENT)) {
            const GumboVector* attributes = &node->v.element.attributes;
            if (likely(has_class_fast(attributes, class_cstr, class_len))) {
                results.emplace_back(create_node_fast(node, tree));
            }
        }
    });
    
    return results;
}

std::shared_ptr<Node> UltraFastWalker::find_by_id_ultra_fast(
    GumboNode* root,
    const Tree* tree,
    const std::string& id
) {
    if (unlikely(!root)) return nullptr;
    
    const char* id_cstr = id.c_str();
    const size_t id_len = id.length();
    
    std::shared_ptr<Node> result;
    
    traverse_depth_first_optimized(root, [&](GumboNode* node) -> bool {
        if (likely(node && node->type == GUMBO_NODE_ELEMENT)) {
            const char* node_id = find_attribute_fast(&node->v.element.attributes, "id");
            if (likely(node_id)) {
                if (config_.enable_simd && id_len >= 8) {
                    if (simd_string_equals(node_id, id_cstr, id_len)) {
                        result = create_node_fast(node, tree);
                        return false;  // Early termination
                    }
                } else {
                    if (strlen(node_id) == id_len && memcmp(node_id, id_cstr, id_len) == 0) {
                        result = create_node_fast(node, tree);
                        return false;  // Early termination
                    }
                }
            }
        }
        return true;  // Continue searching
    });
    
    return result;
}

void UltraFastWalker::collect_all_elements_fast(
    GumboNode* root,
    std::vector<GumboNode*>& out_elements
) {
    if (unlikely(!root)) return;
    
    traverse_depth_first_optimized(root, [&](GumboNode* node) {
        if (likely(node && node->type == GUMBO_NODE_ELEMENT)) {
            out_elements.push_back(node);
        }
    });
}

void UltraFastWalker::collect_all_elements_with_tag_fast(
    GumboNode* root,
    GumboTag target_tag,
    std::vector<GumboNode*>& out_elements
) {
    if (unlikely(!root)) return;
    
    traverse_depth_first_optimized(root, [&](GumboNode* node) {
        if (likely(node && node->type == GUMBO_NODE_ELEMENT)) {
            if (likely(node->v.element.tag == target_tag)) {
                out_elements.push_back(node);
            }
        }
    });
}

std::vector<std::shared_ptr<Node>> UltraFastWalker::find_by_tag_simd(
    GumboNode* root,
    const Tree* tree,
    const char* tag_name,
    size_t tag_len
) {
    std::vector<std::shared_ptr<Node>> results;
    results.reserve(config_.reserve_capacity);
    
    if (unlikely(!root || !tag_name)) return results;
    
    // For very short tags, use the faster enum comparison if possible
    if (tag_len <= 16) {
        GumboTag target_tag = gumbo_tag_enum(tag_name);
        if (target_tag != GUMBO_TAG_UNKNOWN) {
            collect_all_elements_with_tag_fast(root, target_tag, 
                reinterpret_cast<std::vector<GumboNode*>&>(results));
            // Convert to Node objects
            for (auto& gumbo_node : reinterpret_cast<std::vector<GumboNode*>&>(results)) {
                // This is a hack - in reality we'd need a different approach
                // but this shows the concept
            }
            return results;
        }
    }
    
    // SIMD-optimized string matching for custom tags
    traverse_depth_first_optimized(root, [&](GumboNode* node) {
        if (likely(node && node->type == GUMBO_NODE_ELEMENT)) {
            const char* node_name = gumbo_normalized_tagname(node->v.element.tag);
            if (likely(node_name && simd_string_equals(node_name, tag_name, tag_len))) {
                results.emplace_back(create_node_fast(node, tree));
            }
        }
    });
    
    return results;
}

} // namespace cutysoup
