#include "pch.h"
#include "ComponentChildrenManager.h"
#include <sstream>
#include <algorithm>
#include <iomanip>

// Main interface method
ComponentChildrenManager::ComponentChildrenResult ComponentChildrenManager::GetComponentFirstLevelChildren(const ComponentChildrenRequest& request) {
    ComponentChildrenResult result;

    if (request.component_path.empty()) {
        result.error_message = "Component path is required";
        return result;
    }

    SessionInfo sessionInfo = GetSessionInfo();
    if (!sessionInfo.is_valid) {
        result.error_message = "Creo session not available";
        return result;
    }

    try {
        pfcModel_ptr current_model = sessionInfo.session->GetCurrentModel();
        if (!current_model) {
            result.error_message = "No current model loaded";
            return result;
        }

        if (current_model->GetType() != pfcMDL_ASSEMBLY) {
            result.error_message = "Current model is not an assembly";
            return result;
        }

        wfcWAssembly_ptr currentAssembly = wfcWAssembly::cast(current_model);
        if (!currentAssembly) {
            result.error_message = "Failed to cast current model to assembly";
            return result;
        }

        // Check if component_path refers to the current top-level assembly
        std::string currentModelName = "";
        try {
            xstring filename_xstr = current_model->GetFileName();
            currentModelName = XStringToString(filename_xstr);
        } catch (...) {
            // If can't get filename, try origin
            try {
                xstring origin = current_model->GetOrigin();
                std::string origin_str = XStringToString(origin);
                size_t pos = origin_str.find_last_of("/\\");
                if (pos != std::string::npos) {
                    currentModelName = origin_str.substr(pos + 1);
                } else {
                    currentModelName = origin_str;
                }
            } catch (...) {
                currentModelName = "";
            }
        }

        // Parse the component path
        std::string actualPath = request.component_path;
        ParsedPath targetPath;

        if (actualPath.empty()) {
            // Empty path is valid (represents top-level query)
            targetPath.is_valid = true;
        } else {
            targetPath = ParseComponentPath(actualPath);
            if (!targetPath.is_valid) {
                result.error_message = "Invalid component path format";
                return result;
            }

            // If path starts with current model name, skip the first segment
            if (!targetPath.path_segments.empty() &&
                !currentModelName.empty() &&
                CaseInsensitiveCompare(targetPath.path_segments[0], currentModelName)) {
                // Remove first segment as it refers to current assembly itself
                targetPath.path_segments.erase(targetPath.path_segments.begin());
                if (!targetPath.path_segments.empty()) {
                    targetPath.target_component = targetPath.path_segments.back();
                } else {
                    // Only top-level name provided, treat as top-level query
                    actualPath = "";
                }
            }
        }

        // Handle top-level assembly case (empty path after conversion)
        ComponentMatch match;
        bool found = false;

        if (actualPath.empty() || targetPath.path_segments.empty()) {
            // Direct query of top-level assembly, no need to search
            match.found = true;
            match.owner_assembly = currentAssembly;
            found = true;
        } else {
            // Search for the component in the assembly
            found = FindComponentByPath(currentAssembly, targetPath, match);
            if (!found || !match.found) {
                result.error_message = "Component not found at path: " + request.component_path;
                return result;
            }
        }

        // Validate component ID if provided
        if (request.component_id != -1) {
            int actualFeatureId = -1;
            try {
                if (match.feature) {
                    actualFeatureId = match.feature->GetId();
                }
            } catch (...) {
                actualFeatureId = -1;
            }

            if (actualFeatureId != request.component_id) {
                result.error_message = "Component Feature ID mismatch. Expected: " + std::to_string(request.component_id) +
                                     ", Found: " + std::to_string(actualFeatureId);
                return result;
            }
        }

        // Set parent component info
        result.parent_component_path = request.component_path;

        pfcModel_ptr componentModel = nullptr;
        wfcWAssembly_ptr componentAssembly = nullptr;

        if (actualPath.empty()) {
            // Top-level assembly case
            result.parent_component_id = -1;  // No Feature ID for top-level
            result.parent_component_filename = currentModelName;
            componentModel = current_model;
            componentAssembly = currentAssembly;
        } else {
            // Sub-component case
            try {
                if (match.feature) {
                    result.parent_component_id = match.feature->GetId();
                }

                auto modelDescr = match.component_feature->GetModelDescr();
                if (modelDescr) {
                    xstring filename_xstr = modelDescr->GetFileName();
                    result.parent_component_filename = XStringToString(filename_xstr);
                }
            } catch (...) {
                result.parent_component_id = -1;
                result.parent_component_filename = "unknown";
            }

            // Load the component model to get its children
            componentModel = LoadComponentModel(match.component_feature);
            if (!componentModel) {
                result.error_message = "Failed to load component model";
                return result;
            }

            // Check if the component is an assembly (has children)
            if (componentModel->GetType() != pfcMDL_ASSEMBLY) {
                // Part has no children, return empty list
                result.success = true;
                result.message = "Component is a part (no children)";
                result.children_count = 0;
                return result;
            }

            // Cast to assembly and get first level children
            componentAssembly = wfcWAssembly::cast(componentModel);
        }
        if (!componentAssembly) {
            result.error_message = "Failed to cast component to assembly";
            return result;
        }

        // Get all component features at the first level
        pfcFeatures_ptr features = componentAssembly->ListFeaturesByType(xfalse, pfcFEATTYPE_COMPONENT);
        if (!features) {
            result.success = true;
            result.message = "Assembly has no components";
            result.children_count = 0;
            return result;
        }

        int features_count = features->getarraysize();
        if (features_count <= 0) {
            result.success = true;
            result.message = "Assembly has no components";
            result.children_count = 0;
            return result;
        }

        // Collect all first-level children
        for (int i = 0; i < features_count; i++) {
            try {
                pfcFeature_ptr feature = features->get(i);
                if (!feature) continue;

                pfcComponentFeat_ptr compFeat = pfcComponentFeat::cast(feature);
                if (!compFeat) continue;

                ChildComponentInfo childInfo = CreateChildComponentFromFeature(compFeat, 1, request.component_path);
                result.children.push_back(childInfo);

            } catch (...) {
                continue;
            }
        }

        result.children_count = result.children.size();
        result.success = true;
        result.message = "Successfully retrieved " + std::to_string(result.children_count) + " child components";

    } catch (const std::exception& e) {
        result.error_message = "Exception: " + std::string(e.what());
    } catch (...) {
        result.error_message = "Unknown exception occurred";
    }

    return result;
}

// Helper method implementations
ComponentChildrenManager::SessionInfo ComponentChildrenManager::GetSessionInfo() {
    SessionInfo info;
    try {
        info.session = pfcGetCurrentSessionWithCompatibility(pfcC4Compatible);
        info.is_valid = (info.session != nullptr);
    } catch (...) {
        info.is_valid = false;
    }
    return info;
}

ComponentChildrenManager::ParsedPath ComponentChildrenManager::ParseComponentPath(const std::string& full_path) {
    ParsedPath result;

    if (full_path.empty()) {
        return result;
    }

    // Split path by '/' or '\'
    std::string path = full_path;
    std::replace(path.begin(), path.end(), '\\', '/');

    std::istringstream iss(path);
    std::string segment;

    while (std::getline(iss, segment, '/')) {
        if (!segment.empty()) {
            result.path_segments.push_back(segment);
        }
    }

    if (!result.path_segments.empty()) {
        result.target_component = result.path_segments.back();
        result.is_valid = true;
    }

    return result;
}

bool ComponentChildrenManager::FindComponentByPath(wfcWAssembly_ptr assembly,
                                                 const ParsedPath& target_path,
                                                 ComponentMatch& match) {
    if (!assembly || !target_path.is_valid) return false;

    try {
        return RecursiveSearchComponent(assembly, target_path, "", 0, match);
    } catch (...) {
        return false;
    }
}

bool ComponentChildrenManager::RecursiveSearchComponent(wfcWAssembly_ptr currentAssembly,
                                                      const ParsedPath& target_path,
                                                      const std::string& pathSoFar,
                                                      int currentLevel,
                                                      ComponentMatch& match) {
    if (!currentAssembly || currentLevel >= (int)target_path.path_segments.size()) {
        return false;
    }

    try {
        std::string currentTarget = target_path.path_segments[currentLevel];

        pfcFeatures_ptr features = currentAssembly->ListFeaturesByType(xfalse, pfcFEATTYPE_COMPONENT);
        if (!features) return false;

        int features_count = features->getarraysize();

        for (int i = 0; i < features_count; i++) {
            try {
                pfcFeature_ptr feature = features->get(i);
                if (!feature) continue;

                pfcComponentFeat_ptr compFeat = pfcComponentFeat::cast(feature);
                if (!compFeat) continue;

                auto modelDescr = compFeat->GetModelDescr();
                if (!modelDescr) continue;

                xstring filename_xstr = modelDescr->GetFileName();
                std::string componentName = XStringToString(filename_xstr);

                std::string newPath = pathSoFar.empty() ? componentName : pathSoFar + "/" + componentName;

                if (CaseInsensitiveCompare(componentName, currentTarget)) {
                    // Found the target at this level
                    if (currentLevel == (int)target_path.path_segments.size() - 1) {
                        // This is the final target
                        match.feature = feature;
                        match.component_feature = compFeat;
                        match.owner_assembly = currentAssembly;
                        match.actual_path = newPath;
                        match.found = true;
                        return true;
                    } else {
                        // Need to go deeper
                        pfcModel_ptr childModel = LoadComponentModel(compFeat);
                        if (childModel && childModel->GetType() == pfcMDL_ASSEMBLY) {
                            wfcWAssembly_ptr childAssembly = wfcWAssembly::cast(childModel);
                            if (childAssembly) {
                                if (RecursiveSearchComponent(childAssembly, target_path, newPath, currentLevel + 1, match)) {
                                    return true;
                                }
                            }
                        }
                    }
                }

            } catch (...) {
                continue;
            }
        }

    } catch (...) {
        return false;
    }

    return false;
}

ComponentChildrenManager::ChildComponentInfo ComponentChildrenManager::CreateChildComponentFromFeature(pfcComponentFeat_ptr compFeat,
                                                                                                        int level,
                                                                                                        const std::string& parentPath) {
    ChildComponentInfo component;
    component.level = level;
    component.children_count = 0;
    component.is_visible = true;
    component.file_size = "0.0MB";
    component.id = -1; // Default Feature ID

    try {
        // Get Feature ID first
        pfcFeature_ptr feature = pfcFeature::cast(compFeat);
        if (feature) {
            component.id = feature->GetId();
        }

        auto modelDescr = compFeat->GetModelDescr();
        if (modelDescr) {
            xstring filename_xstr = modelDescr->GetFileName();
            component.filename = XStringToString(filename_xstr);

            component.name = component.filename;
            size_t ext_pos = component.name.find_last_of(".");
            if (ext_pos != std::string::npos) {
                component.name = component.name.substr(0, ext_pos);
            }

            if (!component.name.empty()) {
                component.name[0] = std::toupper(component.name[0]);
            }

            // Build paths
            if (parentPath.empty()) {
                component.path = component.filename;
                component.full_path = component.filename;
            } else {
                component.path = parentPath + "/" + component.filename;
                component.full_path = component.path;
            }

            // Determine component type
            std::string ext = "";
            if (ext_pos != std::string::npos) {
                ext = component.filename.substr(ext_pos + 1);
                std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
            }

            if (ext == "asm") {
                component.type = "assembly";
                component.model_type = "MDL_ASSEMBLY";
            } else {
                component.type = "part";
                component.model_type = "MDL_PART";
            }

            // Try to load model and get additional info
            pfcModel_ptr childModel = LoadComponentModel(compFeat);
            if (childModel) {
                component.file_size = GetModelFileSize(childModel);

                if (childModel->GetType() == pfcMDL_ASSEMBLY) {
                    wfcWAssembly_ptr childAssembly = wfcWAssembly::cast(childModel);
                    if (childAssembly) {
                        try {
                            pfcFeatures_ptr childFeatures = childAssembly->ListFeaturesByType(xfalse, pfcFEATTYPE_COMPONENT);
                            if (childFeatures) {
                                component.children_count = childFeatures->getarraysize();
                            }
                        } catch (...) {
                            component.children_count = 0;
                        }
                    }
                }
            }
        }

    } catch (...) {
        component.id = -1;
        component.filename = "unknown_component";
        component.name = "Unknown";
        component.type = "unknown";
        component.model_type = "UNKNOWN";
    }

    return component;
}

pfcModel_ptr ComponentChildrenManager::LoadComponentModel(pfcComponentFeat_ptr compFeat) {
    if (!compFeat) return nullptr;

    try {
        auto modelDescr = compFeat->GetModelDescr();
        if (!modelDescr) return nullptr;

        SessionInfo sessionInfo = GetSessionInfo();
        if (!sessionInfo.is_valid) return nullptr;

        pfcModel_ptr model = sessionInfo.session->GetModelFromDescr(modelDescr);
        if (!model) {
            model = sessionInfo.session->RetrieveModel(modelDescr);
        }

        return model;
    } catch (...) {
        return nullptr;
    }
}

std::string ComponentChildrenManager::GetModelFileSize(pfcModel_ptr model) {
    try {
        return "N/A";
    } catch (...) {
        return "0.0MB";
    }
}

std::string ComponentChildrenManager::XStringToString(const xstring& xstr) {
    try {
        const wchar_t* wstr = xstr;
        if (!wstr) return "";

        int bufferSize = WideCharToMultiByte(CP_UTF8, 0, wstr, -1, nullptr, 0, nullptr, nullptr);
        if (bufferSize <= 0) return "";

        std::string result(bufferSize - 1, '\0');
        WideCharToMultiByte(CP_UTF8, 0, wstr, -1, &result[0], bufferSize, nullptr, nullptr);

        return result;
    } catch (...) {
        return "";
    }
}

xstring ComponentChildrenManager::StringToXString(const std::string& str) {
    try {
        int bufferSize = MultiByteToWideChar(CP_UTF8, 0, str.c_str(), -1, nullptr, 0);
        if (bufferSize <= 0) return xstring();

        std::vector<wchar_t> wstr(bufferSize);
        MultiByteToWideChar(CP_UTF8, 0, str.c_str(), -1, &wstr[0], bufferSize);

        return xstring(&wstr[0]);
    } catch (...) {
        return xstring();
    }
}

bool ComponentChildrenManager::CaseInsensitiveCompare(const std::string& str1, const std::string& str2) {
    if (str1.length() != str2.length()) return false;

    return std::equal(str1.begin(), str1.end(), str2.begin(),
                     [](char a, char b) {
                         return std::tolower(a) == std::tolower(b);
                     });
}