#include "test_common.h"
#include "bs_roformer/inference.h"
#include <cstring>

// We need to test the static helper OR the pipeline.
// Check if Inference::ProcessOverlapAdd is still available and public.
// Yes, it is in inference.h and implemented in inference.cpp.

int main(int argc, char* argv[]) {
    std::cout << "Test: Chunking Logic (Overlap-Add) Verification" << std::endl;
    
    std::string data_dir = GetTestDataDir();
    // Use files generated by export_chunking_debug.py if available, or skip
    // If we included them in package_test_data.py, they might be in activations/ ??
    // No, export_chunking_debug.py puts them in root or specified dir.
    // If packaged, we might have them? 
    // Let's assume they are in data_dir (which might be "golden" root).
    // Note: chunk_in.npy and chunk_out.npy are NOT in 'activations/' subdir normally
    // but in tests/ or root.
    // Let's try loading from data_dir directly.
    
    // Fallback: If not found, try generating? Or just skip?
    // Better: Assume they are present.
    
    // We use load_npy directly as they might not be in activations/
    // We use load_npy directly as they are in tests/ directory
    // Use data_dir (from BSR_TEST_DATA_DIR or default)
    std::string in_path = data_dir + "/chunk_in.npy";
    std::string out_path = data_dir + "/chunk_out.npy";

    if (argc > 1) in_path = argv[1];
    if (argc > 2) out_path = argv[2];
    
    auto [in_ptr, in_shape] = utils::load_npy(in_path);
    if (!in_ptr) {
         // Try checking if it's in the 'activations' subdir (legacy/alternative structure)
         std::string alt_in = data_dir + "/activations/chunk_in.npy";
         auto res = utils::load_npy(alt_in);
         if (res.first) {
             in_ptr = res.first; in_shape = res.second;
             in_path = alt_in;
         }
    }
    
    if (!in_ptr) {
        // Just print absolute path hint for debugging
        std::cout << "[SKIP] chunk_in.npy not found in " << data_dir << " or " << in_path << std::endl;
        return 0; 
    }
    
    auto [out_ptr, out_shape] = utils::load_npy(out_path);
    if (!out_ptr) {
         std::string alt_out = data_dir + "/activations/chunk_out.npy";
         auto res = utils::load_npy(alt_out);
         if (res.first) {
             out_ptr = res.first; out_shape = res.second;
         }
    }
    
    if (!out_ptr) {
         std::cout << "[SKIP] chunk_out.npy not found" << std::endl;
         utils::free_npy_data(in_ptr);
         return 0;
    }
    
    std::vector<float> input_vec(utils::shape_nelements(in_shape));
    std::memcpy(input_vec.data(), in_ptr, input_vec.size()*sizeof(float));
    utils::free_npy_data(in_ptr);
    
    // Expected
    std::vector<float> expected_vec(utils::shape_nelements(out_shape));
    std::memcpy(expected_vec.data(), out_ptr, expected_vec.size()*sizeof(float));
    utils::free_npy_data(out_ptr);
    
    // Run Logic
    int chunk_size = 352800;
    int num_overlap = 2;
    
    std::cout << "  Input size: " << input_vec.size() << std::endl;
    
    // Identity Model
    auto identity = [](const std::vector<float>& chunk) { return std::vector<std::vector<float>>{chunk}; };
    
    // We test the STATIC legacy method because we can't easily mock the pipeline
    // inside Inference class without refactoring it to accept an abstract Model interface.
    
    auto actual_stems = Inference::ProcessOverlapAdd(input_vec, chunk_size, num_overlap, identity);
    std::vector<float> actual = actual_stems[0];
    
    bool pass = CompareAndReport("OverlapAdd Logic", 
                                  expected_vec.data(), expected_vec.size(),
                                  actual.data(), actual.size(),
                                  1e-4f, 1e-4f);
                                  
    if (pass) LOG_PASS(); else LOG_FAIL();
    return pass ? 0 : 1;
}