#include "bulk_operations.h"
#include "tree.h"
#include <algorithm>
#include <cstring>
#include <chrono>
#include <sstream>

namespace cutysoup {

BulkOperations::BulkOperations(const Tree* tree) 
    : tree_(tree), indices_built_(false) {
}

void BulkOperations::build_indices(GumboNode* root) {
    if (!root || indices_built_) return;
    
    auto start_time = std::chrono::high_resolution_clock::now();
    
    // Clear existing indices
    tag_index_.clear();
    class_index_.clear();
    id_index_.clear();
    attribute_index_.clear();
    text_index_.clear();
    all_elements_.clear();
    stats_ = IndexStats{};
    
    // Traverse the entire tree and build indices
    std::vector<GumboNode*> stack;
    stack.reserve(1000);  // Pre-allocate to avoid reallocations
    stack.push_back(root);
    
    while (!stack.empty()) {
        GumboNode* current = stack.back();
        stack.pop_back();
        
        if (!current) continue;
        
        // Index this element if it's an element node
        if (current->type == GUMBO_NODE_ELEMENT) {
            index_element(current);
            all_elements_.push_back(current);  // Add to all elements for nth operations
            stats_.total_elements++;
        }
        
        // Add children to stack for processing
        if (current->type == GUMBO_NODE_ELEMENT) {
            const GumboVector* children = &current->v.element.children;
            for (int i = static_cast<int>(children->length) - 1; i >= 0; --i) {
                stack.push_back(static_cast<GumboNode*>(children->data[i]));
            }
        }
    }
    
    auto end_time = std::chrono::high_resolution_clock::now();
    index_build_time_ = std::chrono::duration<double>(end_time - start_time).count();
    
    // Update statistics
    stats_.unique_tags = tag_index_.size();
    stats_.unique_classes = class_index_.size();
    stats_.unique_ids = id_index_.size();
    stats_.index_memory_bytes = estimate_memory_usage();
    
    indices_built_ = true;
}

void BulkOperations::index_element(GumboNode* node) {
    if (!node || node->type != GUMBO_NODE_ELEMENT) return;
    
    // Index by tag name
    const char* tag_name = gumbo_normalized_tagname(node->v.element.tag);
    if (tag_name) {
        tag_index_[std::string(tag_name)].push_back(node);
    }
    
    // Index by attributes (class, id, and all attr=value pairs)
    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 || !attr->name || !attr->value) continue;
        
        // Index all attribute=value pairs for advanced searching
        std::string attr_key = std::string(attr->name) + "=" + std::string(attr->value);
        attribute_index_[attr_key].push_back(node);
        
        if (strcmp(attr->name, "class") == 0) {
            extract_classes(attr->value, node);
        } else if (strcmp(attr->name, "id") == 0) {
            // IDs should be unique, but we'll take the first one
            std::string id_str(attr->value);
            if (id_index_.find(id_str) == id_index_.end()) {
                id_index_[id_str] = node;
            }
        }
    }
    
    // Index text content (only immediate text children for better performance)
    const GumboVector* children = &node->v.element.children;
    for (unsigned int i = 0; i < children->length; ++i) {
        GumboNode* child = static_cast<GumboNode*>(children->data[i]);
        if (child && child->type == GUMBO_NODE_TEXT && child->v.text.text) {
            std::string text_content = child->v.text.text;
            
            // Trim whitespace for better indexing
            size_t start = text_content.find_first_not_of(" \t\n\r");
            if (start != std::string::npos) {
                size_t end = text_content.find_last_not_of(" \t\n\r");
                text_content = text_content.substr(start, (end - start + 1));
                
                if (!text_content.empty()) {
                    text_index_[text_content].push_back(node);
                }
            }
        }
    }
}

void BulkOperations::extract_classes(const char* class_attr, GumboNode* node) {
    if (!class_attr || !node) return;
    
    const char* pos = class_attr;
    const char* end = class_attr + strlen(class_attr);
    
    while (pos < end) {
        // Skip whitespace
        while (pos < end && (*pos == ' ' || *pos == '\t' || *pos == '\n' || *pos == '\r')) {
            ++pos;
        }
        
        if (pos >= end) break;
        
        // Find end of current class
        const char* class_start = pos;
        while (pos < end && *pos != ' ' && *pos != '\t' && *pos != '\n' && *pos != '\r') {
            ++pos;
        }
        
        // Add this class to the index
        if (pos > class_start) {
            std::string class_name(class_start, pos - class_start);
            class_index_[class_name].push_back(node);
        }
    }
}

std::vector<std::shared_ptr<Node>> BulkOperations::create_nodes_from_raw(
    const std::vector<GumboNode*>& raw_nodes
) const {
    std::vector<std::shared_ptr<Node>> result;
    result.reserve(raw_nodes.size());
    
    for (GumboNode* raw_node : raw_nodes) {
        if (raw_node) {
            if (tree_) {
                result.emplace_back(tree_->get_or_create_node(raw_node));
            } else {
                result.emplace_back(std::make_shared<Node>(raw_node, tree_));
            }
        }
    }
    
    return result;
}

std::vector<std::shared_ptr<Node>> BulkOperations::find_all_by_tag_bulk(const std::string& tag_name) {
    if (!indices_built_) return {};
    
    auto it = tag_index_.find(tag_name);
    if (it == tag_index_.end()) {
        return {};
    }
    
    return create_nodes_from_raw(it->second);
}

std::vector<std::shared_ptr<Node>> BulkOperations::find_all_by_class_bulk(const std::string& class_name) {
    if (!indices_built_) return {};
    
    auto it = class_index_.find(class_name);
    if (it == class_index_.end()) {
        return {};
    }
    
    return create_nodes_from_raw(it->second);
}

std::shared_ptr<Node> BulkOperations::find_by_id_bulk(const std::string& id) {
    if (!indices_built_) return nullptr;
    
    auto it = id_index_.find(id);
    if (it == id_index_.end()) {
        return nullptr;
    }
    
    if (tree_) {
        return tree_->get_or_create_node(it->second);
    }
    return std::make_shared<Node>(it->second, tree_);
}

std::vector<std::shared_ptr<Node>> BulkOperations::find_all_bulk(
    const std::string& tag_name,
    const std::map<std::string, std::string>& attributes
) {
    if (!indices_built_) return {};
    
    std::vector<GumboNode*> candidate_nodes;
    
    // Start with the most selective filter
    if (!tag_name.empty()) {
        // Filter by tag name first (usually most selective)
        auto tag_it = tag_index_.find(tag_name);
        if (tag_it == tag_index_.end()) {
            return {}; // No matching tags
        }
        candidate_nodes = tag_it->second;
    } else if (attributes.size() == 1) {
        // Single attribute search - use indices if available
        auto attr_it = attributes.begin();
        if (attr_it->first == "class") {
            auto class_it = class_index_.find(attr_it->second);
            if (class_it == class_index_.end()) {
                return {}; // No matching class
            }
            candidate_nodes = class_it->second;
        } else if (attr_it->first == "id") {
            auto id_it = id_index_.find(attr_it->second);
            if (id_it == id_index_.end()) {
                return {}; // No matching id
            }
            candidate_nodes = {id_it->second};
        }
    }
    
    // If we don't have good candidate nodes, fall back to full search
    if (candidate_nodes.empty() && !attributes.empty()) {
        // Need to search all elements - collect from tag index
        for (const auto& tag_pair : tag_index_) {
            candidate_nodes.insert(candidate_nodes.end(), 
                                 tag_pair.second.begin(), 
                                 tag_pair.second.end());
        }
    }
    
    // Filter candidates by attributes
    std::vector<GumboNode*> filtered_nodes;
    filtered_nodes.reserve(candidate_nodes.size());
    
    for (GumboNode* node : candidate_nodes) {
        if (!node || node->type != GUMBO_NODE_ELEMENT) continue;
        
        // Check tag name match
        if (!tag_name.empty()) {
            const char* node_tag = gumbo_normalized_tagname(node->v.element.tag);
            if (!node_tag || tag_name != std::string(node_tag)) {
                continue;
            }
        }
        
        // Check attribute matches
        bool attr_match = true;
        const GumboVector* node_attributes = &node->v.element.attributes;
        
        for (const auto& required_attr : attributes) {
            bool found_match = false;
            
            for (unsigned int i = 0; i < node_attributes->length; ++i) {
                GumboAttribute* attr = static_cast<GumboAttribute*>(node_attributes->data[i]);
                if (!attr || !attr->name || !attr->value) continue;
                
                if (required_attr.first == std::string(attr->name)) {
                    if (required_attr.first == "class") {
                        // Special handling for class attributes (space-separated)
                        std::string class_value(attr->value);
                        if (class_value.find(required_attr.second) != std::string::npos) {
                            // More precise class matching
                            std::string search_class = " " + required_attr.second + " ";
                            std::string attr_classes = " " + class_value + " ";
                            found_match = (attr_classes.find(search_class) != std::string::npos);
                        }
                    } else {
                        // Exact match for other attributes
                        found_match = (required_attr.second == std::string(attr->value));
                    }
                    break;
                }
            }
            
            if (!found_match) {
                attr_match = false;
                break;
            }
        }
        
        if (attr_match) {
            filtered_nodes.push_back(node);
        }
    }
    
    return create_nodes_from_raw(filtered_nodes);
}

std::shared_ptr<Node> BulkOperations::find_bulk(
    const std::string& tag_name,
    const std::map<std::string, std::string>& attributes
) {
    auto results = find_all_bulk(tag_name, attributes);
    return results.empty() ? nullptr : results[0];
}

std::vector<std::vector<std::shared_ptr<Node>>> BulkOperations::multi_find_all(
    const std::vector<std::string>& tag_names
) {
    std::vector<std::vector<std::shared_ptr<Node>>> results;
    results.reserve(tag_names.size());
    
    for (const std::string& tag_name : tag_names) {
        results.push_back(find_all_by_tag_bulk(tag_name));
    }
    
    return results;
}

size_t BulkOperations::count_by_tag(const std::string& tag_name) {
    if (!indices_built_) return 0;
    
    auto it = tag_index_.find(tag_name);
    return (it != tag_index_.end()) ? it->second.size() : 0;
}

size_t BulkOperations::count_by_class(const std::string& class_name) {
    if (!indices_built_) return 0;
    
    auto it = class_index_.find(class_name);
    return (it != class_index_.end()) ? it->second.size() : 0;
}

BulkOperations::IndexStats BulkOperations::get_stats() const {
    return stats_;
}

size_t BulkOperations::estimate_memory_usage() const {
    size_t total = 0;
    
    // Tag index memory
    for (const auto& pair : tag_index_) {
        total += pair.first.size();  // Key string
        total += pair.second.size() * sizeof(GumboNode*);  // Vector of pointers
    }
    
    // Class index memory
    for (const auto& pair : class_index_) {
        total += pair.first.size();  // Key string
        total += pair.second.size() * sizeof(GumboNode*);  // Vector of pointers
    }
    
    // ID index memory
    for (const auto& pair : id_index_) {
        total += pair.first.size();  // Key string
        total += sizeof(GumboNode*);  // Single pointer
    }
    
    return total;
}

//==============================================================================
// Advanced Indexing Methods
//==============================================================================

std::vector<std::shared_ptr<Node>> BulkOperations::find_by_attribute(const std::string& attr_name, const std::string& value) {
    if (!indices_built_) return {};
    
    std::string attr_key = attr_name + "=" + value;
    auto it = attribute_index_.find(attr_key);
    if (it != attribute_index_.end()) {
        return create_nodes_from_raw(it->second);
    }
    
    return {};
}

std::vector<std::shared_ptr<Node>> BulkOperations::find_by_text_content(const std::string& text) {
    if (!indices_built_) return {};
    
    auto it = text_index_.find(text);
    if (it != text_index_.end()) {
        return create_nodes_from_raw(it->second);
    }
    
    return {};
}

std::shared_ptr<Node> BulkOperations::get_nth_element(size_t index) {
    if (!indices_built_ || index >= all_elements_.size()) {
        return nullptr;
    }
    
    GumboNode* node = all_elements_[index];
    if (node) {
        if (tree_) {
            return tree_->get_or_create_node(node);
        }
        return std::make_shared<Node>(node, tree_);
    }
    
    return nullptr;
}

} // namespace cutysoup
