#include "pch.h"
#include "ModelSearchEngine.h"
#include <sstream>
#include <iostream>
#include <iomanip>
#include <cctype>
#include <algorithm>

// Singleton implementation
ModelSearchEngine& ModelSearchEngine::Instance() {
    static ModelSearchEngine instance;
    return instance;
}

// Main search interface
ModelSearchResult ModelSearchEngine::SearchModels(const ModelSearchRequest& request) {
    ModelSearchResult result;
    
    // Start timing
    StartTimer();
    
    try {
        // Validate request parameters
        if (request.query.empty()) {
            result.error_message = "Search query cannot be empty";
            return result;
        }
        
        // Clear previous search state
        processed_models_.clear();
        std::vector<SearchResultItem> all_results;
        
        // Choose search method based on scope
        if (request.search_scope == "session" || request.search_scope == "all") {
            // Search from root assemblies to build complete hierarchy paths
            SearchFromRootAssemblies(request, all_results);
        }
        else if (request.search_scope == "current_model") {
            SearchCurrentModel(request, all_results);
        }
        else {
            result.error_message = "Invalid search scope: " + request.search_scope;
            return result;
        }
        
        // Sort and filter results
        SortAndFilterResults(all_results, request);
        
        // Limit result count
        result.total_found = static_cast<int>(all_results.size());
        if (request.max_results > 0 && all_results.size() > static_cast<size_t>(request.max_results)) {
            all_results.resize(request.max_results);
        }
        
        result.results = all_results;
        result.returned_count = static_cast<int>(result.results.size());
        result.success = true;
        result.stats.search_scope_used = request.search_scope;
        
    }
    catch (const pfcXToolkitError& e) {
        result.error_message = "OTK Error occurred during search";
    }
    catch (const std::exception& e) {
        result.error_message = "Search error: " + std::string(e.what());
    }
    catch (...) {
        result.error_message = "Unknown error during model search";
    }
    
    // Set search time
    result.search_time_ms = GetElapsedTimeMs();
    
    return result;
}

// Search models in current session
void ModelSearchEngine::SearchSessionModels(const ModelSearchRequest& request, std::vector<SearchResultItem>& results) {
    try {
        pfcSession_ptr session = pfcGetCurrentSessionWithCompatibility(pfcC4Compatible);
        if (!session) {
            return;
        }
        
        // Get all models in session
        pfcModels_ptr models = session->ListModels();
        if (!models) {
            return;
        }
        
        // Iterate through all models
        for (int i = 0; i < models->getarraysize(); i++) {
            pfcModel_ptr model = models->get(i);
            if (!model) continue;
            
            // Get model information
            std::string model_type = GetModelTypeString(model);
            
            // Get file name for consistent path building
            std::string model_name;
            try {
                xstring filename_xstr = model->GetFileName();
                std::wstring wfilename = filename_xstr;
                model_name = std::string(wfilename.begin(), wfilename.end());
            } catch (...) {
                model_name = GetModelDisplayName(model); // fallback
            }
            
            std::string file_path = GetModelFilePath(model);
            std::string full_path = GetModelFullPath(model);
            
            // Check model type filter
            if (!ShouldIncludeModelType(model_type, request.model_types)) {
                continue;
            }
            
            // Calculate match score
            double match_score = CalculateMatchScore(request, model_name, "name");
            
            // Check similarity threshold
            if (match_score < request.similarity_threshold) {
                continue;
            }
            
            // Create search result item
            SearchResultItem item;
            item.model_name = model_name;
            item.full_path = model_name;  // 对于会话中的根模型，路径就是模型名
            item.file_path = file_path;   // 物理文件路径
            item.model_type = model_type;
            item.match_score = match_score;
            item.match_type = "name";
            item.is_in_session = true;
            item.is_assembly = (model_type == "ASSEMBLY");
            
            // Build match reason
            if (request.match_mode == "prefix") {
                item.match_reason = "Name starts with '" + request.query + "'";
            } else if (request.match_mode == "contains") {
                item.match_reason = "Name contains '" + request.query + "'";
            } else {
                item.match_reason = "Fuzzy match similarity: " + std::to_string(static_cast<int>(match_score * 100)) + "%";
            }
            
            item.matched_keywords.push_back(request.query);
            results.push_back(item);
            
            // If assembly and need to search components
            if (request.include_components && item.is_assembly) {
                try {
                    pfcAssembly_ptr assembly = pfcAssembly::cast(model);
                    if (assembly) {
                        std::set<std::string> local_processed;
                        SearchAssemblyComponents(request, assembly, results, local_processed, model_name);
                    }
                }
                catch (...) {
                    // Assembly search failure doesn't affect main results
                }
            }
        }
        
    }
    catch (...) {
        // Session search failure doesn't throw exception
    }
}

// Search from root assemblies to build complete hierarchy paths
void ModelSearchEngine::SearchFromRootAssemblies(const ModelSearchRequest& request, std::vector<SearchResultItem>& results) {
    try {
        pfcSession_ptr session = pfcGetCurrentSessionWithCompatibility(pfcC4Compatible);
        if (!session) return;
        
        pfcModels_ptr models = session->ListModels();
        if (!models) return;
        
        // Global deduplication map: model_file_path -> longest_hierarchy_path
        std::map<std::string, SearchResultItem> best_results;
        
        // Find root assemblies and search recursively
        for (int i = 0; i < models->getarraysize(); i++) {
            pfcModel_ptr model = models->get(i);
            if (!model) continue;
            
            // Skip if not assembly
            if (model->GetType() != pfcMDL_ASSEMBLY) continue;
            
            // Get model file name for root path
            std::string root_name;
            try {
                xstring filename_xstr = model->GetFileName();
                std::wstring wfilename = filename_xstr;
                root_name = std::string(wfilename.begin(), wfilename.end());
            } catch (...) {
                continue;
            }
            
            // Check if root assembly itself matches
            double match_score = CalculateMatchScore(request, root_name, "name");
            if (match_score >= request.similarity_threshold) {
                SearchResultItem item;
                item.model_name = root_name;
                item.full_path = root_name; // Root assembly path
                item.file_path = GetModelFilePath(model);
                item.model_type = "ASSEMBLY";
                item.match_score = match_score;
                item.match_type = "name";
                item.is_in_session = true;
                item.is_assembly = true;
                item.match_reason = (request.match_mode == "contains") ? 
                    ("Name contains '" + request.query + "'") : 
                    ("Name matches '" + request.query + "'");
                item.matched_keywords.push_back(request.query);
                
                // Add to best results with deduplication
                std::string key = item.file_path;
                if (best_results.find(key) == best_results.end() || 
                    best_results[key].full_path.length() < item.full_path.length()) {
                    best_results[key] = item;
                }
            }
            
            // Search components recursively with complete path
            if (request.include_components) {
                pfcAssembly_ptr assembly = pfcAssembly::cast(model);
                if (assembly) {
                    std::set<std::string> local_processed;
                    SearchAssemblyComponentsWithDedup(request, assembly, best_results, local_processed, root_name);
                }
            }
        }
        
        // Add best results to final results
        for (auto& pair : best_results) {
            results.push_back(pair.second);
        }
        
    } catch (...) {
        // Error in root assembly search
    }
}

// Search current model
void ModelSearchEngine::SearchCurrentModel(const ModelSearchRequest& request, std::vector<SearchResultItem>& results) {
    try {
        pfcSession_ptr session = pfcGetCurrentSessionWithCompatibility(pfcC4Compatible);
        if (!session) {
            return;
        }
        
        pfcModel_ptr current_model = session->GetCurrentModel();
        if (!current_model) {
            return;
        }
        
        // Check current model itself
        // Get file name for consistent path building
        std::string model_name;
        try {
            xstring filename_xstr = current_model->GetFileName();
            std::wstring wfilename = filename_xstr;
            model_name = std::string(wfilename.begin(), wfilename.end());
        } catch (...) {
            model_name = GetModelDisplayName(current_model); // fallback
        }
        
        double match_score = CalculateMatchScore(request, model_name, "name");
        
        if (match_score >= request.similarity_threshold) {
            std::string file_path = GetModelFilePath(current_model);
            std::string full_path = GetModelFullPath(current_model);
            
            SearchResultItem item;
            item.model_name = model_name;
            item.full_path = model_name;  // 当前模型的路径就是模型名
            item.file_path = file_path;   // 物理文件路径
            item.model_type = GetModelTypeString(current_model);
            item.match_score = match_score;
            item.match_type = "name";
            item.is_in_session = true;
            item.is_assembly = (item.model_type == "ASSEMBLY");
            item.match_reason = "Current model match";
            item.matched_keywords.push_back(request.query);
            results.push_back(item);
        }
        
        // If assembly, search its components
        if (request.include_components) {
            try {
                pfcAssembly_ptr assembly = pfcAssembly::cast(current_model);
                if (assembly) {
                    std::set<std::string> local_processed;
                    SearchAssemblyComponents(request, assembly, results, local_processed, model_name);
                }
            }
            catch (...) {
                // Component search failure doesn't affect main results
            }
        }
        
    }
    catch (...) {
        // Current model search failure doesn't throw exception
    }
}

// Search assembly components (recursive)
void ModelSearchEngine::SearchAssemblyComponents(const ModelSearchRequest& request, 
                                               pfcAssembly_ptr assembly, 
                                               std::vector<SearchResultItem>& results,
                                               std::set<std::string>& processed_models,
                                               const std::string& parent_path) {
    if (!assembly) return;
    
    try {
        // Get all component features (including suppressed ones)
        pfcFeatures_ptr features = assembly->ListFeaturesByType(xfalse, pfcFeatureType::pfcFEATTYPE_COMPONENT);
        if (!features) return;
        
        if (features->getarraysize() == 0) return;
        
        for (int i = 0; i < features->getarraysize(); i++) {
            pfcFeature_ptr feature = features->get(i);
            if (!feature) continue;
            
            try {
                pfcComponentFeat_ptr comp_feat = pfcComponentFeat::cast(feature);
                if (!comp_feat) continue;
                
                pfcModelDescriptor_ptr desc = comp_feat->GetModelDescr();
                if (!desc) continue;
                
                pfcModel_ptr comp_model = nullptr;
                try {
                    pfcSession_ptr session = pfcGetCurrentSessionWithCompatibility(pfcC4Compatible);
                    if (session) {
                        comp_model = session->GetModelFromDescr(desc);
                    }
                } catch (pfcXToolkitError&) {
                    // Model not loaded, skip this component
                    continue;
                }
                if (!comp_model) continue;
                
                // Get file name for component path (like hierarchy analysis)
                std::string comp_name;
                try {
                    xstring filename_xstr = desc->GetFileName();
                    std::wstring wfilename = filename_xstr;
                    comp_name = std::string(wfilename.begin(), wfilename.end());
                } catch (...) {
                    comp_name = GetModelDisplayName(comp_model); // fallback
                }
                
                std::string comp_file = GetModelFilePath(comp_model);
                std::string full_path = GetModelFullPath(comp_model);
                
                // Prevent duplicate processing of same model
                if (processed_models.find(comp_file) != processed_models.end()) {
                    continue;
                }
                processed_models.insert(comp_file);
                
                // Check model type filter
                std::string model_type = GetModelTypeString(comp_model);
                if (!ShouldIncludeModelType(model_type, request.model_types)) {
                    continue;
                }
                
                // Calculate match score
                double match_score = CalculateMatchScore(request, comp_name, "component");
                
                // Build tree path for model hierarchy
                std::string tree_path;
                if (parent_path.empty()) {
                    tree_path = comp_name;
                } else {
                    tree_path = parent_path + "/" + comp_name;
                }
                
                if (match_score >= request.similarity_threshold) {
                    SearchResultItem item;
                    item.model_name = comp_name;
                    item.full_path = tree_path;         // Model tree hierarchy path
                    item.file_path = comp_file;         // Physical file path
                    item.model_type = model_type;
                    item.match_score = match_score;
                    item.match_type = "component";
                    item.is_in_session = true;
                    item.is_assembly = (model_type == "ASSEMBLY");
                    item.match_reason = "Component name match in assembly";
                    item.matched_keywords.push_back(request.query);
                    results.push_back(item);
                }
                
                // If this component is also an assembly, search recursively
                if (model_type == "ASSEMBLY") {
                    try {
                        pfcAssembly_ptr sub_assembly = pfcAssembly::cast(comp_model);
                        if (sub_assembly) {
                            SearchAssemblyComponents(request, sub_assembly, results, processed_models, tree_path);
                        }
                    }
                    catch (...) {
                        // Sub-assembly search failure doesn't affect other results
                    }
                }
                
            }
            catch (...) {
                // Single component processing failure doesn't affect other components
                continue;
            }
        }
        
    }
    catch (...) {
        // Assembly component search failure doesn't throw exception
    }
}

// Search assembly components with global deduplication (recursive)
void ModelSearchEngine::SearchAssemblyComponentsWithDedup(const ModelSearchRequest& request, 
                                                         pfcAssembly_ptr assembly, 
                                                         std::map<std::string, SearchResultItem>& best_results,
                                                         std::set<std::string>& processed_models,
                                                         const std::string& parent_path) {
    if (!assembly) return;
    
    try {
        // Get all component features (including suppressed ones)
        pfcFeatures_ptr features = assembly->ListFeaturesByType(xfalse, pfcFeatureType::pfcFEATTYPE_COMPONENT);
        if (!features) return;
        
        if (features->getarraysize() == 0) return;
        
        for (int i = 0; i < features->getarraysize(); i++) {
            pfcFeature_ptr feature = features->get(i);
            if (!feature) continue;
            
            try {
                pfcComponentFeat_ptr comp_feat = pfcComponentFeat::cast(feature);
                if (!comp_feat) continue;
                
                pfcModelDescriptor_ptr desc = comp_feat->GetModelDescr();
                if (!desc) continue;
                
                pfcModel_ptr comp_model = nullptr;
                try {
                    pfcSession_ptr session = pfcGetCurrentSessionWithCompatibility(pfcC4Compatible);
                    if (session) {
                        comp_model = session->GetModelFromDescr(desc);
                    }
                } catch (pfcXToolkitError&) {
                    // Model not loaded, skip this component
                    continue;
                }
                if (!comp_model) continue;
                
                // Get file name for component path (like hierarchy analysis)
                std::string comp_name;
                try {
                    xstring filename_xstr = desc->GetFileName();
                    std::wstring wfilename = filename_xstr;
                    comp_name = std::string(wfilename.begin(), wfilename.end());
                } catch (...) {
                    comp_name = GetModelDisplayName(comp_model); // fallback
                }
                
                std::string comp_file = GetModelFilePath(comp_model);
                std::string full_path = GetModelFullPath(comp_model);
                
                // Prevent duplicate processing of same model
                if (processed_models.find(comp_file) != processed_models.end()) {
                    continue;
                }
                processed_models.insert(comp_file);
                
                // Check model type filter
                std::string model_type = GetModelTypeString(comp_model);
                if (!ShouldIncludeModelType(model_type, request.model_types)) {
                    continue;
                }
                
                // Calculate match score
                double match_score = CalculateMatchScore(request, comp_name, "component");
                
                // Build tree path for model hierarchy
                std::string tree_path;
                if (parent_path.empty()) {
                    tree_path = comp_name;
                } else {
                    tree_path = parent_path + "/" + comp_name;
                }
                
                if (match_score >= request.similarity_threshold) {
                    SearchResultItem item;
                    item.model_name = comp_name;
                    item.full_path = tree_path;         // Model tree hierarchy path
                    item.file_path = comp_file;         // Physical file path
                    item.model_type = model_type;
                    item.match_score = match_score;
                    item.match_type = "component";
                    item.is_in_session = true;
                    item.is_assembly = (model_type == "ASSEMBLY");
                    item.match_reason = "Component name match in assembly";
                    item.matched_keywords.push_back(request.query);
                    
                    // Add to best results with deduplication by file path
                    // Keep the result with longest hierarchy path
                    std::string key = comp_file;
                    if (best_results.find(key) == best_results.end() || 
                        best_results[key].full_path.length() < item.full_path.length()) {
                        best_results[key] = item;
                    }
                }
                
                // If this component is also an assembly, search recursively
                if (model_type == "ASSEMBLY") {
                    try {
                        pfcAssembly_ptr sub_assembly = pfcAssembly::cast(comp_model);
                        if (sub_assembly) {
                            SearchAssemblyComponentsWithDedup(request, sub_assembly, best_results, processed_models, tree_path);
                        }
                    }
                    catch (...) {
                        // Sub-assembly search failure doesn't affect other results
                    }
                }
                
            }
            catch (...) {
                // Single component processing failure doesn't affect other components
                continue;
            }
        }
        
    }
    catch (...) {
        // Assembly component search failure doesn't throw exception
    }
}

// Calculate match score
double ModelSearchEngine::CalculateMatchScore(const ModelSearchRequest& request, 
                                           const std::string& target_name,
                                           const std::string& match_type) {
    if (request.match_mode == "prefix") {
        return MatchesPrefix(ToLowerCase(target_name), ToLowerCase(request.query)) ? 1.0 : 0.0;
    }
    else if (request.match_mode == "contains") {
        return MatchesContains(ToLowerCase(target_name), ToLowerCase(request.query)) ? 0.8 : 0.0;
    }
    else if (request.match_mode == "fuzzy") {
        double score = CalculateFuzzyMatch(ToLowerCase(request.query), ToLowerCase(target_name));
        
        // Adjust score based on match type
        if (match_type == "name") {
            score *= 1.0;  // Name match no adjustment
        } else if (match_type == "component") {
            score *= 0.9;  // Component match slightly lower score
        }
        
        return score;
    }
    
    return 0.0;
}

// Sort and filter results
void ModelSearchEngine::SortAndFilterResults(std::vector<SearchResultItem>& results, 
                                          const ModelSearchRequest& request) {
    // Sort by match score in descending order
    std::sort(results.begin(), results.end(), 
        [](const SearchResultItem& a, const SearchResultItem& b) {
            if (std::abs(a.match_score - b.match_score) < 0.001) {
                // Same score, prioritize name matches
                if (a.match_type != b.match_type) {
                    if (a.match_type == "name") return true;
                    if (b.match_type == "name") return false;
                }
                // Then sort by name alphabetically
                return a.model_name < b.model_name;
            }
            return a.match_score > b.match_score;
        });
}

// Fuzzy matching algorithm implementation
double ModelSearchEngine::CalculateFuzzyMatch(const std::string& query, const std::string& target) {
    if (query.empty() || target.empty()) {
        return 0.0;
    }
    
    // Exact match
    if (query == target) {
        return 1.0;
    }
    
    // Prefix match
    if (target.find(query) == 0) {
        return 0.95;
    }
    
    // Contains match
    if (target.find(query) != std::string::npos) {
        return 0.8;
    }
    
    // Edit distance match
    double edit_distance = CalculateEditDistance(query, target);
    size_t max_len = (query.length() > target.length()) ? query.length() : target.length();
    
    if (max_len == 0) return 1.0;
    
    double similarity = 1.0 - (edit_distance / max_len);
    return (similarity > 0.0) ? similarity : 0.0;
}

// String matching tools implementation
bool ModelSearchEngine::MatchesPrefix(const std::string& text, const std::string& prefix) {
    return text.size() >= prefix.size() && text.substr(0, prefix.size()) == prefix;
}

bool ModelSearchEngine::MatchesContains(const std::string& text, const std::string& substring) {
    return text.find(substring) != std::string::npos;
}

double ModelSearchEngine::CalculateEditDistance(const std::string& s1, const std::string& s2) {
    const size_t len1 = s1.size();
    const size_t len2 = s2.size();
    
    std::vector<std::vector<int>> dp(len1 + 1, std::vector<int>(len2 + 1));
    
    for (size_t i = 0; i <= len1; ++i) dp[i][0] = static_cast<int>(i);
    for (size_t j = 0; j <= len2; ++j) dp[0][j] = static_cast<int>(j);
    
    for (size_t i = 1; i <= len1; ++i) {
        for (size_t j = 1; j <= len2; ++j) {
            if (s1[i-1] == s2[j-1]) {
                dp[i][j] = dp[i-1][j-1];
            } else {
                int a = dp[i-1][j];
                int b = dp[i][j-1]; 
                int c = dp[i-1][j-1];
                dp[i][j] = 1 + ((a < b) ? ((a < c) ? a : c) : ((b < c) ? b : c));
            }
        }
    }
    
    return static_cast<double>(dp[len1][len2]);
}

// Helper methods implementation
std::string ModelSearchEngine::GetModelTypeString(pfcModel_ptr model) {
    if (!model) return "UNKNOWN";
    
    pfcModelType type = model->GetType();
    switch (type) {
        case pfcModelType::pfcMDL_PART: return "PART";
        case pfcModelType::pfcMDL_ASSEMBLY: return "ASSEMBLY";
        case pfcModelType::pfcMDL_DRAWING: return "DRAWING";
        default: return "OTHER";
    }
}

std::string ModelSearchEngine::GetModelDisplayName(pfcModel_ptr model) {
    if (!model) return "";
    
    try {
        xstring name = model->GetInstanceName();
        std::wstring wname = name;
        return std::string(wname.begin(), wname.end());
    }
    catch (...) {
        try {
            xstring filename = model->GetOrigin();
            std::wstring wfilename = filename;
            std::string full_name(wfilename.begin(), wfilename.end());
            
            // Extract filename part
            size_t last_slash = full_name.find_last_of("\\/");
            if (last_slash != std::string::npos) {
                return full_name.substr(last_slash + 1);
            }
            return full_name;
        }
        catch (...) {
            return "Unknown";
        }
    }
}

std::string ModelSearchEngine::GetModelFilePath(pfcModel_ptr model) {
    if (!model) return "";
    
    try {
        xstring origin = model->GetOrigin();
        std::wstring worigin = origin;
        return std::string(worigin.begin(), worigin.end());
    }
    catch (...) {
        return "";
    }
}

std::string ModelSearchEngine::GetModelFullPath(pfcModel_ptr model) {
    if (!model) return "";
    
    try {
        // Try GetFullName first (includes full path)
        xstring full_name = model->GetFullName();
        std::wstring wfull_name = full_name;
        return std::string(wfull_name.begin(), wfull_name.end());
    }
    catch (...) {
        // Fallback to GetOrigin if GetFullName fails
        try {
            xstring origin = model->GetOrigin();
            std::wstring worigin = origin;
            return std::string(worigin.begin(), worigin.end());
        }
        catch (...) {
            return "";
        }
    }
}

bool ModelSearchEngine::ShouldIncludeModelType(const std::string& model_type, 
                                             const std::vector<std::string>& type_filters) {
    if (type_filters.empty()) {
        return true;  // No filter, include all types
    }
    
    for (const auto& filter : type_filters) {
        if (filter == "MDL_PART" && model_type == "PART") return true;
        if (filter == "MDL_ASSEMBLY" && model_type == "ASSEMBLY") return true;
        if (filter == "MDL_DRAWING" && model_type == "DRAWING") return true;
        if (filter == model_type) return true;
    }
    
    return false;
}

// String processing tools implementation
std::string ModelSearchEngine::ToLowerCase(const std::string& str) {
    std::string result = str;
    std::transform(result.begin(), result.end(), result.begin(), ::tolower);
    return result;
}

std::vector<std::string> ModelSearchEngine::SplitString(const std::string& str, char delimiter) {
    std::vector<std::string> tokens;
    std::stringstream ss(str);
    std::string token;
    
    while (std::getline(ss, token, delimiter)) {
        tokens.push_back(TrimString(token));
    }
    
    return tokens;
}

std::string ModelSearchEngine::TrimString(const std::string& str) {
    size_t start = str.find_first_not_of(" \t\n\r");
    if (start == std::string::npos) return "";
    
    size_t end = str.find_last_not_of(" \t\n\r");
    return str.substr(start, end - start + 1);
}

std::string ModelSearchEngine::BuildComponentPath(const std::string& parent_path, const std::string& component_name) {
    if (parent_path.empty()) {
        return component_name;
    }
    return parent_path + " > " + component_name;
}

// Performance timing methods
void ModelSearchEngine::StartTimer() {
    search_start_time_ = std::chrono::high_resolution_clock::now();
}

std::string ModelSearchEngine::GetElapsedTimeMs() {
    auto end_time = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - search_start_time_);
    return std::to_string(duration.count()) + "ms";
}