#include "model_detector.h"
#include <iostream>
#include <vector>
#include <filesystem>

namespace fs = std::filesystem;

// Test function specifically for Qwen model detection
void testQwenDetection(const std::string& modelPath) {
    std::cout << "\n=== Testing Qwen Model Detection for: " << modelPath << " ===" << std::endl;

    if (!fs::exists(modelPath)) {
        std::cout << "ERROR: Model file does not exist!" << std::endl;
        return;
    }

    try {
        // Test ModelDetector
        ModelDetectionResult result = ModelDetector::detectModel(modelPath);

        std::cout << "📋 Detection Results:" << std::endl;
        std::cout << "  Architecture: " << result.architectureName << std::endl;
        std::cout << "  Architecture Enum: " << static_cast<int>(result.architecture) << std::endl;
        
        // Check if it's correctly detected as QWEN2VL
        bool isCorrectlyDetected = (result.architecture == ModelArchitecture::QWEN2VL);
        std::cout << "  ✅ Correctly detected as QWEN2VL: " << (isCorrectlyDetected ? "YES" : "NO") << std::endl;
        
        if (!isCorrectlyDetected) {
            std::cout << "  ❌ FAILED: Expected QWEN2VL but got " << result.architectureName << std::endl;
            std::cout << "  💡 This indicates the fix is not working properly" << std::endl;
        } else {
            std::cout << "  🎉 SUCCESS: Qwen model correctly detected!" << std::endl;
        }

        // Show tensor names for debugging
        std::cout << "\n🔍 Tensor Analysis (first 10 tensors):" << std::endl;
        int count = 0;
        for (const auto& tensorName : result.tensorNames) {
            if (count >= 10) break;
            std::cout << "  " << (count + 1) << ". " << tensorName << std::endl;
            count++;
        }
        
        if (result.tensorNames.size() > 10) {
            std::cout << "  ... and " << (result.tensorNames.size() - 10) << " more tensors" << std::endl;
        }

        // Show metadata
        std::cout << "\n📄 Metadata:" << std::endl;
        for (const auto& [key, value] : result.metadata) {
            std::cout << "  " << key << ": " << value << std::endl;
        }

        // Show suggested parameters
        if (!result.suggestedParams.empty()) {
            std::cout << "\n⚙️  Suggested Parameters:" << std::endl;
            for (const auto& [key, value] : result.suggestedParams) {
                std::cout << "  " << key << ": " << value << std::endl;
            }
        }

    } catch (const std::exception& e) {
        std::cout << "❌ Error during model detection: " << e.what() << std::endl;
    }
}

int main() {
    std::cout << "🧪 Qwen Model Detection Test Suite" << std::endl;
    std::cout << "===================================" << std::endl;

    // Test with available Qwen models
    std::vector<std::string> qwenModelPaths = {
        "/data/SD_MODELS/diffusion_models/Qwen-Image-Edit-2509-Q3_K_S.gguf",
        "/data/SD_MODELS/diffusion_models/Qwen-Image-Pruning-13b-Q4_0.gguf",
        "/data/SD_MODELS/diffusion_models/qwen-image-Q2_K.gguf"
    };

    int successCount = 0;
    int totalTests = 0;

    for (const auto& modelPath : qwenModelPaths) {
        if (fs::exists(modelPath)) {
            totalTests++;
            testQwenDetection(modelPath);
            
            // Check if detection was successful
            ModelDetectionResult result = ModelDetector::detectModel(modelPath);
            if (result.architecture == ModelArchitecture::QWEN2VL) {
                successCount++;
            }
        } else {
            std::cout << "\n⚠️  Skipping test for non-existent model: " << modelPath << std::endl;
        }
    }

    // Summary
    std::cout << "\n🎯 Test Summary:" << std::endl;
    std::cout << "  Total Qwen models tested: " << totalTests << std::endl;
    std::cout << "  Successfully detected as QWEN2VL: " << successCount << std::endl;
    std::cout << "  Success rate: " << (totalTests > 0 ? (successCount * 100 / totalTests) : 0) << "%" << std::endl;

    if (successCount == totalTests && totalTests > 0) {
        std::cout << "  🎉 ALL TESTS PASSED! Qwen detection fix is working correctly." << std::endl;
    } else if (totalTests > 0) {
        std::cout << "  ❌ Some tests failed. The fix may need adjustment." << std::endl;
    } else {
        std::cout << "  ⚠️  No Qwen models found to test." << std::endl;
    }

    std::cout << "\n🏁 Qwen Detection Test Complete!" << std::endl;

    return (successCount == totalTests) ? 0 : 1;
}