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

namespace cutysoup {

ZeroContextBatch::ZeroContextBatch(const Tree* tree) 
    : tree_(tree), cache_enabled_(true) {
    build_fast_indices();
}

void ZeroContextBatch::build_fast_indices() {
    if (!tree_ || !tree_->get_output() || !tree_->get_output()->root) return;
    
    // Build comprehensive indices for zero-context-switch operations
    all_elements_.clear();
    tag_indices_.clear();
    class_indices_.clear();
    id_indices_.clear();
    
    // Traverse once and build all indices
    std::vector<GumboNode*> stack;
    stack.reserve(1000);
    stack.push_back(tree_->get_output()->root);
    
    while (!stack.empty()) {
        GumboNode* current = stack.back();
        stack.pop_back();
        
        if (!current) continue;
        
        if (current->type == GUMBO_NODE_ELEMENT) {
            size_t element_index = all_elements_.size();
            all_elements_.push_back(current);
            
            // Index by tag
            const char* tag_name = gumbo_normalized_tagname(current->v.element.tag);
            if (tag_name) {
                tag_indices_[std::string(tag_name)].push_back(element_index);
            }
            
            // Index by attributes (class and id)
            const GumboVector* attributes = &current->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;
                
                if (strcmp(attr->name, "class") == 0) {
                    // Extract individual classes
                    std::vector<std::string> classes = extract_classes(attr->value);
                    for (const std::string& class_name : classes) {
                        class_indices_[class_name].push_back(element_index);
                    }
                } else if (strcmp(attr->name, "id") == 0) {
                    id_indices_[std::string(attr->value)] = element_index;
                }
            }
        }
        
        // Add children to stack
        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]));
            }
        }
    }
}

std::vector<BatchResult> ZeroContextBatch::execute_batch(const std::vector<BatchQuery>& queries) {
    auto start_time = std::chrono::high_resolution_clock::now();
    
    std::vector<BatchResult> results;
    results.reserve(queries.size());
    
    for (const BatchQuery& query : queries) {
        std::string cache_key = generate_cache_key(query);
        BatchResult result;
        
        // Try cache first
        if (cache_enabled_ && try_get_cached_result(cache_key, result)) {
            stats_.cache_hits++;
            results.push_back(result);
            continue;
        }
        
        // Execute operation
        switch (query.operation) {
            case BatchOp::FIND_ALL_TAG: {
                auto it = tag_indices_.find(query.target);
                if (it != tag_indices_.end()) {
                    std::vector<ElementInfo> elements;
                    elements.reserve(it->second.size());
                    for (size_t idx : it->second) {
                        elements.push_back(extract_element_info(all_elements_[idx], idx));
                    }
                    result = elements;
                } else {
                    result = std::vector<ElementInfo>{};
                }
                break;
            }
            
            case BatchOp::FIND_ALL_CLASS: {
                auto it = class_indices_.find(query.target);
                if (it != class_indices_.end()) {
                    std::vector<ElementInfo> elements;
                    elements.reserve(it->second.size());
                    for (size_t idx : it->second) {
                        elements.push_back(extract_element_info(all_elements_[idx], idx));
                    }
                    result = elements;
                } else {
                    result = std::vector<ElementInfo>{};
                }
                break;
            }
            
            case BatchOp::FIND_BY_ID: {
                auto it = id_indices_.find(query.target);
                if (it != id_indices_.end()) {
                    std::vector<ElementInfo> elements;
                    elements.push_back(extract_element_info(all_elements_[it->second], it->second));
                    result = elements;
                } else {
                    result = std::vector<ElementInfo>{};
                }
                break;
            }
            
            case BatchOp::COUNT_TAG: {
                auto it = tag_indices_.find(query.target);
                result = (it != tag_indices_.end()) ? it->second.size() : size_t(0);
                break;
            }
            
            case BatchOp::COUNT_CLASS: {
                auto it = class_indices_.find(query.target);
                result = (it != class_indices_.end()) ? it->second.size() : size_t(0);
                break;
            }
            
            case BatchOp::EXTRACT_TEXT: {
                std::vector<std::string> texts = get_all_text_by_tag(query.target);
                result = texts;
                break;
            }

            case BatchOp::EXTRACT_ATTRIBUTES: {
                std::vector<std::string> attributes = get_all_attributes(query.target, query.attribute_name);
                result = attributes;
                break;
            }
            
            case BatchOp::GET_TAG_NAMES: {
                std::vector<std::string> tag_names;
                tag_names.reserve(tag_indices_.size());
                for (const auto& pair : tag_indices_) {
                    tag_names.push_back(pair.first);
                }
                result = tag_names;
                break;
            }
        }
        
        // Cache result
        if (cache_enabled_) {
            cache_result(cache_key, result);
        }
        
        results.push_back(result);
        stats_.total_operations++;
    }
    
    auto end_time = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);
    stats_.total_time_saved_ms += duration.count() / 1000.0;
    stats_.context_switches_avoided += queries.size();
    
    return results;
}

std::vector<std::string> ZeroContextBatch::get_all_text_by_tag(const std::string& tag_name) {
    std::vector<std::string> texts;
    
    auto it = tag_indices_.find(tag_name);
    if (it == tag_indices_.end()) return texts;
    
    texts.reserve(it->second.size());
    for (size_t idx : it->second) {
        std::string text = extract_text_content(all_elements_[idx], true);
        if (!text.empty()) {
            texts.push_back(text);
        }
    }
    
    return texts;
}

std::vector<std::string> ZeroContextBatch::get_all_text_by_class(const std::string& class_name) {
    std::vector<std::string> texts;
    
    auto it = class_indices_.find(class_name);
    if (it == class_indices_.end()) return texts;
    
    texts.reserve(it->second.size());
    for (size_t idx : it->second) {
        std::string text = extract_text_content(all_elements_[idx], true);
        if (!text.empty()) {
            texts.push_back(text);
        }
    }
    
    return texts;
}

std::vector<std::string> ZeroContextBatch::get_all_attributes(const std::string& tag_name, const std::string& attr_name) {
    std::vector<std::string> attributes;
    
    auto it = tag_indices_.find(tag_name);
    if (it == tag_indices_.end()) return attributes;
    
    attributes.reserve(it->second.size());
    for (size_t idx : it->second) {
        GumboNode* node = all_elements_[idx];
        const GumboVector* attrs = &node->v.element.attributes;
        
        for (unsigned int i = 0; i < attrs->length; ++i) {
            GumboAttribute* attr = static_cast<GumboAttribute*>(attrs->data[i]);
            if (attr && attr->name && attr->value && attr_name == attr->name) {
                attributes.emplace_back(attr->value);
                break;
            }
        }
    }
    
    return attributes;
}

std::unordered_map<std::string, size_t> ZeroContextBatch::count_all_tags() {
    std::unordered_map<std::string, size_t> counts;
    counts.reserve(tag_indices_.size());
    
    for (const auto& pair : tag_indices_) {
        counts[pair.first] = pair.second.size();
    }
    
    return counts;
}

std::unordered_map<std::string, size_t> ZeroContextBatch::count_all_classes() {
    std::unordered_map<std::string, size_t> counts;
    counts.reserve(class_indices_.size());
    
    for (const auto& pair : class_indices_) {
        counts[pair.first] = pair.second.size();
    }
    
    return counts;
}

void ZeroContextBatch::stream_elements_by_tag(const std::string& tag_name, 
                                             std::function<bool(const ElementInfo&)> callback) {
    auto it = tag_indices_.find(tag_name);
    if (it == tag_indices_.end()) return;
    
    for (size_t idx : it->second) {
        ElementInfo info = extract_element_info(all_elements_[idx], idx);
        if (!callback(info)) {
            break;  // Callback returned false, stop streaming
        }
    }
}

std::vector<std::shared_ptr<Node>> ZeroContextBatch::create_nodes_from_indices(const std::vector<size_t>& indices) {
    std::vector<std::shared_ptr<Node>> nodes;
    nodes.reserve(indices.size());
    
    for (size_t idx : indices) {
        if (idx < all_elements_.size()) {
            if (tree_) {
                nodes.emplace_back(tree_->get_or_create_node(all_elements_[idx]));
            } else {
                nodes.emplace_back(std::make_shared<Node>(all_elements_[idx], tree_));
            }
        }
    }
    
    return nodes;
}

ElementInfo ZeroContextBatch::extract_element_info(GumboNode* node, size_t index) const {
    ElementInfo info;
    info.element_index = index;
    
    if (!node || node->type != GUMBO_NODE_ELEMENT) return info;
    
    // Extract tag name
    const char* tag_name = gumbo_normalized_tagname(node->v.element.tag);
    if (tag_name) {
        info.tag_name = tag_name;
    }
    
    // Extract text content (lightweight)
    info.text_content = extract_text_content(node, false);  // Don't include children for performance
    
    // Extract attributes
    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;
        
        std::string attr_name(attr->name);
        std::string attr_value(attr->value);
        
        info.attributes[attr_name] = attr_value;
        
        if (attr_name == "id") {
            info.id = attr_value;
        } else if (attr_name == "class") {
            info.classes = extract_classes(attr->value);
        }
    }
    
    return info;
}

std::string ZeroContextBatch::extract_text_content(GumboNode* node, bool include_children) const {
    if (!node) return "";
    
    if (node->type == GUMBO_NODE_TEXT) {
        return std::string(node->v.text.text);
    }
    
    if (node->type == GUMBO_NODE_ELEMENT && include_children) {
        std::ostringstream text_stream;
        const GumboVector* children = &node->v.element.children;
        
        for (unsigned int i = 0; i < children->length; ++i) {
            GumboNode* child = static_cast<GumboNode*>(children->data[i]);
            std::string child_text = extract_text_content(child, true);
            if (!child_text.empty()) {
                text_stream << child_text;
            }
        }
        
        return text_stream.str();
    }
    
    return "";
}

std::vector<std::string> ZeroContextBatch::extract_classes(const char* class_attr) const {
    std::vector<std::string> classes;
    if (!class_attr) return classes;
    
    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;
        }
        
        if (pos > class_start) {
            classes.emplace_back(class_start, pos - class_start);
        }
    }
    
    return classes;
}

std::string ZeroContextBatch::generate_cache_key(const BatchQuery& query) const {
    std::ostringstream key;
    key << static_cast<int>(query.operation) << ":" << query.target << ":" << query.attribute_name;
    return key.str();
}

bool ZeroContextBatch::try_get_cached_result(const std::string& key, BatchResult& result) const {
    auto it = result_cache_.find(key);
    if (it != result_cache_.end()) {
        result = it->second;
        return true;
    }
    return false;
}

void ZeroContextBatch::cache_result(const std::string& key, const BatchResult& result) const {
    result_cache_[key] = result;
}

void ZeroContextBatch::clear_cache() {
    result_cache_.clear();
}

} // namespace cutysoup
