#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <chrono>
#include <thread>
#include <mutex>
#include <atomic>
#include <condition_variable>
#include <sys/mman.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <cstring>
#include <simdjson.h>
#include <iomanip>
#include <filesystem>
#include <optional>
#include <algorithm>

// Platform-specific includes for memory info
#if defined(__linux__) || defined(__linux) || defined(linux)
    #include <sys/sysinfo.h>  // For system memory info on Linux
#elif defined(__APPLE__) || defined(__MACH__)
    #include <mach/mach.h>    // For system memory info on macOS
#endif

// Arrow headers
#include <arrow/api.h>
#include <arrow/io/api.h>
#include <arrow/json/api.h>
#include <arrow/ipc/api.h>
#include <arrow/filesystem/api.h>
#include <arrow/buffer.h>
#include <arrow/util/thread_pool.h>
#include <arrow/util/parallel.h>
#include <arrow/memory_pool.h>

// SIMD optimization headers
#if defined(__AVX2__) || defined(HAVE_AVX2)
    #define USE_AVX2
    #include <immintrin.h> // AVX2, AVX, SSE4.2, SSE4.1, SSSE3, SSE3, SSE2, SSE
#elif defined(__SSE4_2__) || defined(HAVE_SSE42)
    #define USE_SSE4_2
    #include <nmmintrin.h> // SSE4.2, SSE4.1, SSSE3, SSE3, SSE2, SSE
#elif defined(__SSE4_1__) || defined(HAVE_SSE41)
    #define USE_SSE4_1
    #include <smmintrin.h> // SSE4.1, SSSE3, SSE3, SSE2, SSE
#elif defined(__SSSE3__) || defined(HAVE_SSSE3)
    #define USE_SSSE3
    #include <tmmintrin.h> // SSSE3, SSE3, SSE2, SSE
#elif defined(__SSE3__) || defined(HAVE_SSE3)
    #define USE_SSE3
    #include <pmmintrin.h> // SSE3, SSE2, SSE
#elif defined(__SSE2__) || defined(HAVE_SSE2)
    #define USE_SSE2
    #include <emmintrin.h> // SSE2, SSE
#endif

// For ARM NEON support
#if defined(__ARM_NEON) || defined(HAVE_NEON)
    #define USE_NEON
    #include <arm_neon.h>
#endif

// Add filesystem namespace for directory operations
namespace fs = std::filesystem;

// Constants for memory management
constexpr size_t SMALL_FILE_THRESHOLD = 256ULL * 1024ULL * 1024ULL;  // 256 MB - use in-memory for these
constexpr size_t LARGE_FILE_THRESHOLD = 2ULL * 1024ULL * 1024ULL * 1024ULL;  // 2 GB - use disk-based for these
constexpr size_t MEMORY_RESERVE_MARGIN = 512ULL * 1024ULL * 1024ULL;  // 512 MB - keep this much free
constexpr size_t DEFAULT_BUFFER_SIZE = 64ULL * 1024ULL * 1024ULL;  // 64 MB default buffer
constexpr size_t LARGE_BUFFER_SIZE = 256ULL * 1024ULL * 1024ULL;  // 256 MB buffer for large operations

// Processing mode enum
enum class ProcessingMode {
    Auto,           // Automatically choose based on file size and memory
    ForceMemory,    // Force in-memory processing
    ForceDisk       // Force disk-based processing
};

// Structure to hold file info for memory mapping
struct MappedFile {
    int fd;
    size_t size;
    void* data;

    MappedFile(const std::string& path) {
        fd = open(path.c_str(), O_RDONLY);
        if (fd == -1) {
            throw std::runtime_error("Failed to open file: " + path + " - " + std::string(strerror(errno)));
        }

        struct stat sb;
        if (fstat(fd, &sb) == -1) {
            close(fd);
            throw std::runtime_error("Failed to get file size: " + std::string(strerror(errno)));
        }

        size = sb.st_size;
        data = mmap(NULL, size, PROT_READ, MAP_PRIVATE, fd, 0);
        if (data == MAP_FAILED) {
            close(fd);
            throw std::runtime_error("Failed to mmap file: " + std::string(strerror(errno)));
        }
    }

    ~MappedFile() {
        if (data != MAP_FAILED) {
            munmap(data, size);
        }
        if (fd != -1) {
            close(fd);
        }
    }
};

// Structure to hold Arrow data in memory
struct ArrowMemoryChunk {
    std::shared_ptr<arrow::Schema> schema;
    std::vector<std::shared_ptr<arrow::RecordBatch>> batches;
};

// Helper function to get available system memory (cross-platform)
size_t get_available_memory() {
#if defined(__linux__) || defined(__linux) || defined(linux)
    // Linux implementation
    struct sysinfo info;
    if (sysinfo(&info) != 0) {
        std::cerr << "Warning: Failed to get system memory info, assuming 4GB" << std::endl;
        return 4ULL * 1024 * 1024 * 1024;  // Default to 4GB if we can't determine
    }
    
    size_t free_ram = info.freeram * info.mem_unit;
    size_t free_swap = info.freeswap * info.mem_unit;
    
    // Use free RAM plus a portion of swap, but leave a safety margin
    size_t usable_memory = free_ram + (free_swap / 2);
    if (usable_memory > MEMORY_RESERVE_MARGIN) {
        usable_memory -= MEMORY_RESERVE_MARGIN;
    }
    
    return usable_memory;
    
#elif defined(__APPLE__) || defined(__MACH__)
    // macOS implementation
    mach_port_t host_port = mach_host_self();
    mach_msg_type_number_t host_size = sizeof(vm_statistics64_data_t) / sizeof(integer_t);
    vm_size_t page_size;
    vm_statistics64_data_t vm_stats;
    
    host_page_size(host_port, &page_size);
    
    if (host_statistics64(host_port, HOST_VM_INFO64, (host_info64_t)&vm_stats, &host_size) != KERN_SUCCESS) {
        std::cerr << "Warning: Failed to get macOS memory info, assuming 4GB" << std::endl;
        return 4ULL * 1024 * 1024 * 1024;  // Default to 4GB if we can't determine
    }
    
    // Calculate free memory (free + inactive)
    size_t free_memory = (vm_stats.free_count + vm_stats.inactive_count) * page_size;
    
    // Leave a safety margin
    if (free_memory > MEMORY_RESERVE_MARGIN) {
        free_memory -= MEMORY_RESERVE_MARGIN;
    }
    
    return free_memory;
    
#else
    // Default implementation for other platforms
    std::cerr << "Warning: System memory detection not implemented for this platform, assuming 4GB" << std::endl;
    return 4ULL * 1024 * 1024 * 1024;  // Default to 4GB
#endif
}

// Function to determine ideal processing mode based on file size and available memory
ProcessingMode determine_processing_mode(size_t file_size, ProcessingMode requested_mode) {
    if (requested_mode == ProcessingMode::ForceMemory) {
        return ProcessingMode::ForceMemory;
    }
    
    if (requested_mode == ProcessingMode::ForceDisk) {
        return ProcessingMode::ForceDisk;
    }
    
    // For Auto mode, determine based on file size and available memory
    size_t available_memory = get_available_memory();
    
    // Small files always use in-memory processing
    if (file_size <= SMALL_FILE_THRESHOLD) {
        return ProcessingMode::ForceMemory;
    }
    
    // Large files use disk-based processing unless we have plenty of memory
    if (file_size >= LARGE_FILE_THRESHOLD) {
        // If we have at least 3x the file size in available memory, we can use in-memory
        if (available_memory >= file_size * 3) {
            std::cout << "Large file detected, but sufficient memory available. Using in-memory processing." << std::endl;
            return ProcessingMode::ForceMemory;
        }
        return ProcessingMode::ForceDisk;
    }
    
    // Medium files use memory-based if we have enough (2x file size)
    if (available_memory >= file_size * 2) {
        return ProcessingMode::ForceMemory;
    }
    
    // Otherwise use disk-based
    return ProcessingMode::ForceDisk;
}

// Function to calculate optimal batch size based on file size and available memory
size_t calculate_optimal_batch_size(size_t file_size, size_t line_count, ProcessingMode mode) {
    // Start with a reasonable default
    size_t target_batch_size = 50000;
    
    // For very large files, use larger batches to reduce the number of batches
    if (file_size > 10ULL * 1024 * 1024 * 1024) {  // 10GB+
        target_batch_size = 500000;
    } else if (file_size > 1ULL * 1024 * 1024 * 1024) {  // 1GB+
        target_batch_size = 100000;
    } else if (file_size < 100 * 1024 * 1024) {  // <100MB
        target_batch_size = 10000;
    }
    
    // In memory mode, we can use larger batches
    if (mode == ProcessingMode::ForceMemory) {
        target_batch_size *= 2;
    }
    
    // Never make batches larger than the total line count
    return std::min(target_batch_size, line_count);
}

// Calculate optimal number of threads based on system and file
size_t calculate_optimal_threads(size_t file_size) {
    // Get available hardware threads
    size_t hardware_threads = std::thread::hardware_concurrency();
    if (hardware_threads == 0) hardware_threads = 4; // Fallback
    
    // For small files, limit thread count to avoid overhead
    if (file_size < 10 * 1024 * 1024) {  // < 10 MB
        return std::max(size_t(1), std::min(size_t(2), hardware_threads));
    }
    
    // For medium files, use a portion of available threads
    if (file_size < 1024 * 1024 * 1024) { // < 1 GB
        return std::max(size_t(2), std::min(size_t(4), hardware_threads));
    }
    
    // For large files, use all threads, but cap at a reasonable maximum to avoid thrashing
    return std::min(hardware_threads, size_t(32));
}

// Debug function to examine JSON content
void debug_json_line(const char* data, size_t start, size_t end, const char* description) {
    std::string line(data + start, end - start);
    std::cout << "\n=== " << description << " ===\n";
    std::cout << "Length: " << line.length() << " bytes\n";
    
    // Print the line with escaped characters
    std::cout << "Content: ";
    for (size_t i = 0; i < line.length(); i++) {
        char c = line[i];
        if (isprint(c)) {
            std::cout << c;
        } else {
            // Print non-printable chars as hex
            printf("\\x%02x", static_cast<unsigned char>(c));
        }
    }
    std::cout << std::endl;
}

// Helper function to validate and potentially fix JSON
bool validate_and_fix_json(const char* data, size_t& start, size_t& end, simdjson::ondemand::parser& parser) {
    // Skip empty lines
    if (end <= start) {
        return false;
    }
    
    // Try parsing as is
    try {
        simdjson::padded_string padded_line(data + start, end - start);
        auto doc = parser.iterate(padded_line);
        return true; // Valid JSON
    } catch (const simdjson::simdjson_error& e) {
        // Try to fix common issues
        
        // 1. Try with trimming trailing whitespace
        size_t trim_end = end;
        while (trim_end > start && std::isspace(data[trim_end - 1])) {
            trim_end--;
        }
        
        if (trim_end < end) {
            try {
                simdjson::padded_string padded_line(data + start, trim_end - start);
                auto doc = parser.iterate(padded_line);
                end = trim_end; // Update the end position
                return true;    // Fixed by trimming whitespace
            } catch (...) {
                // Still invalid, continue trying
            }
        }
        
        // 2. Try with progressive truncation (for incomplete objects)
        size_t test_end = trim_end;
        while (test_end > start) {
            test_end--;
            // Only test at meaningful boundaries
            if (data[test_end] == '}' || data[test_end] == ']') {
                try {
                    simdjson::padded_string padded_line(data + start, test_end + 1 - start);
                    auto doc = parser.iterate(padded_line);
                    end = test_end + 1; // Update the end position
                    return true;    // Fixed by truncation
                } catch (...) {
                    // Still invalid, continue trying
                }
            }
        }
        
        // 3. Check if this is an incomplete object
        // Count opening and closing braces
        int brace_count = 0;
        int bracket_count = 0;
        
        for (size_t i = start; i < end; i++) {
            if (data[i] == '{') brace_count++;
            else if (data[i] == '}') brace_count--;
            else if (data[i] == '[') bracket_count++;
            else if (data[i] == ']') bracket_count--;
        }
        
        // If we have unclosed braces or brackets, try to fix
        if (brace_count > 0 || bracket_count > 0) {
            // Create a repaired JSON with closing braces/brackets
            std::string repaired(data + start, end - start);
            
            // Add closing braces/brackets
            while (brace_count > 0) {
                repaired += '}';
                brace_count--;
            }
            
            while (bracket_count > 0) {
                repaired += ']';
                bracket_count--;
            }
            
            try {
                simdjson::padded_string padded_line(repaired.c_str(), repaired.length());
                auto doc = parser.iterate(padded_line);
                // We can't modify the original data, but we logged that repair is possible
                std::cerr << "Last line could be fixed by adding closing braces/brackets" << std::endl;
                return false; // We can't actually use the repaired version in-place
            } catch (...) {
                // Still invalid
            }
        }
    }
    
    return false; // Couldn't fix
}

// Hexadecimal dump utility function
void dump_hex(const char* data, size_t start, size_t end) {
    std::cout << "Hex dump of line (" << (end-start) << " bytes):" << std::endl;
    
    const size_t BYTES_PER_LINE = 16;
    char ascii[BYTES_PER_LINE + 1];
    ascii[BYTES_PER_LINE] = '\0';
    
    for (size_t i = start; i < end; i += BYTES_PER_LINE) {
        printf("%08lx: ", i - start);
        
        // Print hex values
        for (size_t j = 0; j < BYTES_PER_LINE; j++) {
            if (i + j < end) {
                printf("%02x ", static_cast<unsigned char>(data[i + j]));
                ascii[j] = (isprint(data[i + j])) ? data[i + j] : '.';
            } else {
                printf("   ");
                ascii[j] = ' ';
            }
        }
        
        printf(" %s\n", ascii);
    }
}

// Ultra-optimized parallel line scanner with empty last line fix
std::vector<std::pair<size_t, size_t>> find_line_indices_optimized(const MappedFile& file, int num_threads) {
    const char* data = static_cast<const char*>(file.data);
    const size_t size = file.size;
    
    // Pre-fault memory pages to ensure they're loaded
    const size_t page_size = sysconf(_SC_PAGESIZE);
    #pragma omp parallel for
    for (size_t i = 0; i < size; i += page_size * 16) {
        volatile char c = data[i]; // Touch to ensure page is in memory
    }
    
    // First pass: count newlines per thread chunk
    std::vector<size_t> lines_per_thread(num_threads, 0);
    
    #pragma omp parallel for
    for (int t = 0; t < num_threads; t++) {
        const size_t chunk_size = size / num_threads;
        const size_t start = t * chunk_size;
        const size_t end = (t == num_threads - 1) ? size : (t + 1) * chunk_size;
        
        size_t count = 0;
        
        #if defined(USE_AVX2)
            // AVX2 scanning implementation
            const __m256i newlines = _mm256_set1_epi8('\n');
            const size_t simd_end = start + ((end - start) / 32) * 32;
            
            // Loop unrolling for better throughput (process 128 bytes per iteration)
            size_t i = start;
            for (; i + 128 <= simd_end; i += 128) {
                // Process 4 chunks of 32 bytes
                const __m256i chars1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(data + i));
                const __m256i chars2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(data + i + 32));
                const __m256i chars3 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(data + i + 64));
                const __m256i chars4 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(data + i + 96));
                
                const __m256i matches1 = _mm256_cmpeq_epi8(chars1, newlines);
                const __m256i matches2 = _mm256_cmpeq_epi8(chars2, newlines);
                const __m256i matches3 = _mm256_cmpeq_epi8(chars3, newlines);
                const __m256i matches4 = _mm256_cmpeq_epi8(chars4, newlines);
                
                const unsigned mask1 = _mm256_movemask_epi8(matches1);
                const unsigned mask2 = _mm256_movemask_epi8(matches2);
                const unsigned mask3 = _mm256_movemask_epi8(matches3);
                const unsigned mask4 = _mm256_movemask_epi8(matches4);
                
                count += _mm_popcnt_u32(mask1) + _mm_popcnt_u32(mask2) + 
                         _mm_popcnt_u32(mask3) + _mm_popcnt_u32(mask4);
            }
            
            // Handle remaining 32-byte chunks
            for (; i < simd_end; i += 32) {
                const __m256i chars = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(data + i));
                const __m256i matches = _mm256_cmpeq_epi8(chars, newlines);
                const unsigned mask = _mm256_movemask_epi8(matches);
                count += _mm_popcnt_u32(mask);
            }
            
            // Handle remaining bytes
            for (size_t j = i; j < end; j++) {
                if (data[j] == '\n') {
                    count++;
                }
            }
        #elif defined(USE_SSE4_2) || defined(USE_SSE4_1) || defined(USE_SSSE3) || defined(USE_SSE3) || defined(USE_SSE2)
            // SSE scanning implementation
            const __m128i newlines = _mm_set1_epi8('\n');
            const size_t simd_end = start + ((end - start) / 16) * 16;
            
            // Loop unrolling
            size_t i = start;
            for (; i + 64 <= simd_end; i += 64) {
                // Process 4 chunks of 16 bytes
                const __m128i chars1 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(data + i));
                const __m128i chars2 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(data + i + 16));
                const __m128i chars3 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(data + i + 32));
                const __m128i chars4 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(data + i + 48));
                
                const __m128i matches1 = _mm_cmpeq_epi8(chars1, newlines);
                const __m128i matches2 = _mm_cmpeq_epi8(chars2, newlines);
                const __m128i matches3 = _mm_cmpeq_epi8(chars3, newlines);
                const __m128i matches4 = _mm_cmpeq_epi8(chars4, newlines);
                
                const unsigned mask1 = _mm_movemask_epi8(matches1);
                const unsigned mask2 = _mm_movemask_epi8(matches2);
                const unsigned mask3 = _mm_movemask_epi8(matches3);
                const unsigned mask4 = _mm_movemask_epi8(matches4);
                
                #if defined(USE_SSE4_2) || defined(USE_POPCNT)
                    // Use hardware popcnt if available
                    count += _mm_popcnt_u32(mask1) + _mm_popcnt_u32(mask2) + 
                             _mm_popcnt_u32(mask3) + _mm_popcnt_u32(mask4);
                #else
                    // Fallback popcount implementation
                    auto popcount = [](unsigned x) {
                        x = x - ((x >> 1) & 0x55555555);
                        x = (x & 0x33333333) + ((x >> 2) & 0x33333333);
                        x = (x + (x >> 4)) & 0x0F0F0F0F;
                        x = x + (x >> 8);
                        x = x + (x >> 16);
                        return x & 0x3F;
                    };
                    
                    count += popcount(mask1) + popcount(mask2) + 
                             popcount(mask3) + popcount(mask4);
                #endif
            }
            
            // Handle remaining 16-byte chunks
            for (; i < simd_end; i += 16) {
                const __m128i chars = _mm_loadu_si128(reinterpret_cast<const __m128i*>(data + i));
                const __m128i matches = _mm_cmpeq_epi8(chars, newlines);
                const unsigned mask = _mm_movemask_epi8(matches);
                
                #if defined(USE_SSE4_2) || defined(USE_POPCNT)
                    count += _mm_popcnt_u32(mask);
                #else
                    // Count set bits manually for older CPUs
                    for (unsigned bit = 0; bit < 16; bit++) {
                        if (mask & (1u << bit)) {
                            count++;
                        }
                    }
                #endif
            }
            
            // Handle remaining bytes
            for (size_t j = i; j < end; j++) {
                if (data[j] == '\n') {
                    count++;
                }
            }
        #elif defined(USE_NEON)
            // ARM NEON implementation
            const uint8x16_t newlines = vdupq_n_u8('\n');
            const size_t simd_end = start + ((end - start) / 16) * 16;
            
            for (size_t i = start; i < simd_end; i += 16) {
                const uint8x16_t chars = vld1q_u8(reinterpret_cast<const uint8_t*>(data + i));
                const uint8x16_t matches = vceqq_u8(chars, newlines);
                
                // Count matches (various methods depending on NEON version)
                uint8_t counts[16];
                vst1q_u8(counts, matches);
                
                for (int j = 0; j < 16; j++) {
                    count += (counts[j] != 0);
                }
            }
            
            // Handle remaining bytes
            for (size_t i = simd_end; i < end; i++) {
                if (data[i] == '\n') {
                    count++;
                }
            }
        #else
            // Fallback to use memchr which is highly optimized in standard libraries
            const char* ptr = data + start;
            const char* chunk_end = data + end;
            
            while (ptr < chunk_end) {
                ptr = static_cast<const char*>(memchr(ptr, '\n', chunk_end - ptr));
                if (!ptr) break;
                count++;
                ptr++; // Move past this newline
            }
        #endif
        
        lines_per_thread[t] = count;
    }
    
    // Calculate total lines and pre-allocate result vectors
    size_t total_newlines = 0;
    for (size_t count : lines_per_thread) {
        total_newlines += count;
    }
    
    // Pre-allocate result vector (including first line start at position 0)
    std::vector<size_t> line_starts;
    line_starts.reserve(total_newlines + 1);
    line_starts.push_back(0); // First line always starts at position 0
    
    // Calculate thread offsets for direct placement
    std::vector<size_t> offsets(num_threads);
    size_t offset = 1; // Start at 1 because position 0 is already filled
    for (int t = 0; t < num_threads; t++) {
        offsets[t] = offset;
        offset += lines_per_thread[t];
    }
    
    // Resize line_starts to exact size needed to avoid undefined behavior
    line_starts.resize(offset);
    
    // Second pass: collect line start positions
    #pragma omp parallel for
    for (int t = 0; t < num_threads; t++) {
        const size_t chunk_size = size / num_threads;
        const size_t start = t * chunk_size;
        const size_t end = (t == num_threads - 1) ? size : (t + 1) * chunk_size;
        
        size_t pos = offsets[t]; // Position in output array
        
        #if defined(USE_AVX2)
            const __m256i newlines = _mm256_set1_epi8('\n');
            const size_t simd_end = start + ((end - start) / 32) * 32;
            
            for (size_t i = start; i < simd_end; i += 32) {
                const __m256i chars = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(data + i));
                const __m256i matches = _mm256_cmpeq_epi8(chars, newlines);
                const uint32_t mask = _mm256_movemask_epi8(matches);
                
                if (mask) { // Only process if there are newlines
                    uint32_t remaining = mask;
                    while (remaining) {
                        uint32_t bit_pos = __builtin_ctz(remaining); // Get position of lowest set bit
                        size_t newline_pos = i + bit_pos;
                        if (newline_pos + 1 < size) {
                            line_starts[pos++] = newline_pos + 1;
                        }
                        remaining &= (remaining - 1); // Clear the lowest set bit
                    }
                }
            }
            
            // Handle remaining bytes
            for (size_t i = simd_end; i < end; i++) {
                if (data[i] == '\n' && i + 1 < size) {
                    line_starts[pos++] = i + 1;
                }
            }
        #elif defined(USE_SSE4_2) || defined(USE_SSE4_1) || defined(USE_SSSE3) || defined(USE_SSE3) || defined(USE_SSE2)
            const __m128i newlines = _mm_set1_epi8('\n');
            const size_t simd_end = start + ((end - start) / 16) * 16;
            
            for (size_t i = start; i < simd_end; i += 16) {
                const __m128i chars = _mm_loadu_si128(reinterpret_cast<const __m128i*>(data + i));
                const __m128i matches = _mm_cmpeq_epi8(chars, newlines);
                const uint32_t mask = _mm_movemask_epi8(matches);
                
                if (mask) { // Only process if there are newlines
                    uint32_t remaining = mask;
                    while (remaining) {
                        #if defined(USE_SSE4_2) || defined(__GNUC__) || defined(__clang__)
                            uint32_t bit_pos = __builtin_ctz(remaining); // Get position of lowest set bit
                        #else
                            // Fallback bit scan for MSVC without BMI
                            unsigned long bit_pos;
                            _BitScanForward(&bit_pos, remaining);
                        #endif
                        
                        size_t newline_pos = i + bit_pos;
                        if (newline_pos + 1 < size) {
                            line_starts[pos++] = newline_pos + 1;
                        }
                        remaining &= (remaining - 1); // Clear the lowest set bit
                    }
                }
            }
            
            // Handle remaining bytes
            for (size_t i = simd_end; i < end; i++) {
                if (data[i] == '\n' && i + 1 < size) {
                    line_starts[pos++] = i + 1;
                }
            }
        #elif defined(USE_NEON)
            // ARM NEON implementation
            const uint8x16_t newlines = vdupq_n_u8('\n');
            const size_t simd_end = start + ((end - start) / 16) * 16;
            
            for (size_t i = start; i < simd_end; i += 16) {
                const uint8x16_t chars = vld1q_u8(reinterpret_cast<const uint8_t*>(data + i));
                const uint8x16_t matches = vceqq_u8(chars, newlines);
                
                // Store matches
                uint8_t match_results[16];
                vst1q_u8(match_results, matches);
                
                // Process matches
                for (int j = 0; j < 16; j++) {
                    if (match_results[j]) {
                        size_t newline_pos = i + j;
                        if (newline_pos + 1 < size) {
                            line_starts[pos++] = newline_pos + 1;
                        }
                    }
                }
            }
            
            // Handle remaining bytes
            for (size_t i = simd_end; i < end; i++) {
                if (data[i] == '\n' && i + 1 < size) {
                    line_starts[pos++] = i + 1;
                }
            }
        #else
            // Fallback to use memchr
            const char* ptr = data + start;
            const char* chunk_end = data + end;
            
            while (ptr < chunk_end) {
                ptr = static_cast<const char*>(memchr(ptr, '\n', chunk_end - ptr));
                if (!ptr) break;
                
                size_t newline_pos = ptr - data;
                if (newline_pos + 1 < size) {
                    line_starts[pos++] = newline_pos + 1;
                }
                
                ptr++; // Move past this newline
            }
        #endif
    }
    
    // Create line index pairs
    std::vector<std::pair<size_t, size_t>> indices;
    indices.reserve(line_starts.size());
    
    // Check if file ends with a newline - key for detecting empty trailing line
    bool ends_with_newline = (size > 0 && data[size - 1] == '\n');
    size_t buffer = ends_with_newline ? 2 : 1;
    
    // Build line index pairs
    for (size_t i = 0; i < line_starts.size(); i++) {
        // Skip the last line if it's after the last newline in a file ending with newline
        if (i == line_starts.size() - 1 && ends_with_newline) {
            // This is the line after the last newline in a file ending with newline
            // Skip it entirely as it's an empty line
            continue;
        }
        
        size_t start = line_starts[i];
        size_t end;
        
        if (i < line_starts.size() - buffer) {
            // End is just before the start of the next line
            end = line_starts[i+1] - 1;
            // Adjust for CR if present (CRLF line endings)
            if (end > start && data[end-1] == '\r') {
                end--;
            }
        } else {
            // Last line
            end = size;
            if (end > 0 && data[end-1] == '\n') end--;
            if (end > 0 && data[end-1] == '\r') end--;
        }
        
        // Only add non-empty lines
        if (end >= start) {
            indices.emplace_back(start, end);
        }
    }
    
    return indices;
}

// Optimized schema inference to sample fewer records
arrow::Result<std::shared_ptr<arrow::Schema>> infer_schema_from_sample(
    const char* data,
    const std::vector<std::pair<size_t, size_t>>& indices,
    size_t max_sample_size = 50  // Reduced from 1000 to 50
) {
    // Create a sample JSONL buffer for schema inference
    std::string sample_buffer;
    size_t sample_count = std::min(max_sample_size, indices.size());
    size_t step = indices.size() / sample_count;
    if (step == 0) step = 1;
    
    for (size_t i = 0; i < indices.size(); i += step) {
        if (sample_buffer.size() >= 256 * 1024) break; // Limit to ~256KB (was 1MB)
        
        const auto& [line_start, line_end] = indices[i];
        const size_t line_length = line_end - line_start;
        
        if (line_length > 0) {
            sample_buffer.append(data + line_start, line_length);
            sample_buffer.push_back('\n');
        }
    }
    
    // Use Arrow to infer schema
    auto input = std::make_shared<arrow::io::BufferReader>(
        std::make_shared<arrow::Buffer>(
            reinterpret_cast<const uint8_t*>(sample_buffer.data()),
            sample_buffer.size()
        )
    );
    
    auto read_options = arrow::json::ReadOptions::Defaults();
    read_options.use_threads = false;  // Disable threading in schema inference
    
    auto parse_options_for_schema = arrow::json::ParseOptions::Defaults();
    
    // Create a JSON reader just for schema inference
    ARROW_ASSIGN_OR_RAISE(auto reader, 
        arrow::json::TableReader::Make(arrow::default_memory_pool(), input, read_options, parse_options_for_schema));
    
    // Read a small batch to infer schema
    ARROW_ASSIGN_OR_RAISE(auto table, reader->Read());
    
    return table->schema();
}

// Process chunk and write directly to JSONL format
void process_jsonl_chunk(
    const char* data,
    const std::vector<std::pair<size_t, size_t>>& indices,
    size_t start_idx,
    size_t end_idx,
    std::atomic<size_t>& progress_counter,
    simdjson::ondemand::parser& parser,
    std::mutex& output_mutex,
    std::string& last_json,
    size_t thread_id,
    const std::string& output_dir,
    const std::string& base_filename,
    size_t total_threads
) {
    // Buffer for output using direct file I/O
    constexpr size_t BUFFER_SIZE = LARGE_BUFFER_SIZE; // 256MB buffer
    std::unique_ptr<char[]> buffer(new char[BUFFER_SIZE]);
    size_t buffer_pos = 0;

    // Create output filepath
    std::string output_file = fs::path(output_dir) / fs::path(base_filename + "_" + std::to_string(thread_id) + ".jsonl");
    
    // Open a thread-specific output file
    int fd = open(output_file.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644);
    if (fd == -1) {
        std::cerr << "Thread " << thread_id << ": Failed to open output file: " << output_file 
                 << " - " << strerror(errno) << std::endl;
        return;
    }
    
    // Batch size tracking
    size_t batch_size = 0;
    constexpr size_t BATCH_UPDATE_SIZE = 10000;
    
    size_t total_valid_rows = 0;
    size_t total_invalid_rows = 0;
    
    for (size_t i = start_idx; i < end_idx; ++i) {
        const auto& [line_start, line_end] = indices[i];
        const size_t line_length = line_end - line_start;

        if (line_length > 0) {
            try {
                // For the last line in the file, add extra validation and fixes
                if (i == end_idx - 1 && thread_id == total_threads - 1) {
                    size_t line_start_copy = line_start;
                    size_t line_end_copy = line_end;
                    
                    if (validate_and_fix_json(data, line_start_copy, line_end_copy, parser)) {
                        if (line_end_copy != line_end) {
                            // Use the fixed line
                            // Check if buffer has space, if not flush it
                            if (buffer_pos + (line_end_copy - line_start_copy) + 1 > BUFFER_SIZE) {
                                write(fd, buffer.get(), buffer_pos);
                                buffer_pos = 0;
                            }
                            
                            // Copy fixed line to buffer
                            std::memcpy(buffer.get() + buffer_pos, data + line_start_copy, line_end_copy - line_start_copy);
                            buffer_pos += (line_end_copy - line_start_copy);
                            buffer.get()[buffer_pos++] = '\n';
                            
                            batch_size++;
                            total_valid_rows++;
                            
                            // Save this as the last JSON for display
                            {
                                std::lock_guard<std::mutex> lock(output_mutex);
                                last_json = std::string(data + line_start_copy, line_end_copy - line_start_copy);
                            }
                            
                            continue; // Skip regular processing
                        }
                    } else {
                        debug_json_line(data, line_start, line_end, "Problematic last line");
                    }
                }

                // Use simdjson to parse the line
                simdjson::padded_string padded_line(data + line_start, line_length);
                auto doc = parser.iterate(padded_line);

                // Check if buffer has space, if not flush it
                if (buffer_pos + line_length + 1 > BUFFER_SIZE) {
                    write(fd, buffer.get(), buffer_pos);
                    buffer_pos = 0;
                }

                // Copy line to buffer
                std::memcpy(buffer.get() + buffer_pos, data + line_start, line_length);
                buffer_pos += line_length;
                buffer.get()[buffer_pos++] = '\n';

                // Save this as the last JSON for display (if it's the last record)
                if (i == end_idx - 1 && thread_id == total_threads - 1) {
                    std::lock_guard<std::mutex> lock(output_mutex);
                    last_json = std::string(data + line_start, line_length);
                }
                
                batch_size++;
                total_valid_rows++;
                
                if (batch_size >= BATCH_UPDATE_SIZE) {
                    progress_counter.fetch_add(batch_size, std::memory_order_relaxed);
                    batch_size = 0;
                }
            } catch (const simdjson::simdjson_error& e) {
                // Special handling for the last line in the file
                if (i == end_idx - 1 && thread_id == total_threads - 1) {
                    std::cerr << "Last line was invalid JSON: " << e.what() << std::endl;
                }
                
                total_invalid_rows++;
                // Skip invalid JSON
            } catch (...) {
                total_invalid_rows++;
                // Skip other errors
            }
        }
    }

    // Final flush of buffer
    if (buffer_pos > 0) {
        write(fd, buffer.get(), buffer_pos);
    }
    
    // Close output file
    close(fd);

    // Update progress for any remaining records
    if (batch_size > 0) {
        progress_counter.fetch_add(batch_size, std::memory_order_relaxed);
    }
    
    // Report invalid rows if they exist
    if (total_invalid_rows > 0) {
        std::cerr << "Thread " << thread_id << ": Processed " << (total_valid_rows + total_invalid_rows) 
                 << " lines, " << total_invalid_rows << " invalid JSON objects skipped (" 
                 << (total_invalid_rows * 100.0 / (total_valid_rows + total_invalid_rows)) << "%)" << std::endl;
    }
}

// Optimize batch processing function for Arrow
bool process_and_write_batch(
    const std::string& jsonl_buffer,
    std::shared_ptr<arrow::ipc::RecordBatchWriter>& writer,
    std::shared_ptr<arrow::io::FileOutputStream>& output,
    arrow::MemoryPool* pool,
    std::shared_ptr<arrow::Schema> schema,
    size_t thread_id
) {
    try {
        // Create an input stream from the buffer with zero-copy where possible
        auto buffer_ptr = std::make_shared<arrow::Buffer>(
            reinterpret_cast<const uint8_t*>(jsonl_buffer.data()), 
            jsonl_buffer.size()
        );
        auto input = std::make_shared<arrow::io::BufferReader>(buffer_ptr);
        
        // Configure JSON reader options - disable internal threading to avoid contention
        auto read_options = arrow::json::ReadOptions::Defaults();
        read_options.use_threads = false;  // Don't use threads in JSON reader since we're already parallelized
        read_options.block_size = std::min(static_cast<int64_t>(jsonl_buffer.size()), 
                                         static_cast<int64_t>(8 * 1024 * 1024));  // 8MB blocks or less
        
        // Set parse options with schema if available
        auto parse_options = arrow::json::ParseOptions::Defaults();
        if (schema) {
            parse_options.unexpected_field_behavior = arrow::json::UnexpectedFieldBehavior::Ignore;
        }
        
        // Create JSON reader
        auto reader_result = arrow::json::TableReader::Make(pool, input, read_options, parse_options);
        if (!reader_result.ok()) {
            std::cerr << "Thread " << thread_id << ": Failed to create JSON reader: " 
                     << reader_result.status().ToString() << std::endl;
            return false;
        }
        auto reader = reader_result.ValueOrDie();
        
        // Read the table
        auto table_result = reader->Read();
        if (!table_result.ok()) {
            std::cerr << "Thread " << thread_id << ": Failed to read JSON data: " 
                     << table_result.status().ToString() << std::endl;
            return false;
        }
        auto table = table_result.ValueOrDie();
        
        // Initialize writer if needed (first batch)
        if (!writer) {
            auto options = arrow::ipc::IpcWriteOptions::Defaults();
            options.write_legacy_ipc_format = false;  // Use newer format
            options.max_recursion_depth = 128;  // Increase for complex nested data
            
            auto writer_result = arrow::ipc::MakeFileWriter(output.get(), table->schema(), options);
            if (!writer_result.ok()) {
                std::cerr << "Thread " << thread_id << ": Failed to create Arrow writer: " 
                         << writer_result.status().ToString() << std::endl;
                return false;
            }
            writer = writer_result.ValueOrDie();
        }
        
        // Convert table to batch - we need to manually flatten the chunks
        auto batch_result = table->CombineChunksToBatch();
        if (!batch_result.ok()) {
            std::cerr << "Thread " << thread_id << ": Failed to create record batch: " 
                     << batch_result.status().ToString() << std::endl;
            return false;
        }
        auto batch = batch_result.ValueOrDie();
        
        // Write the batch
        auto write_status = writer->WriteRecordBatch(*batch);
        if (!write_status.ok()) {
            std::cerr << "Thread " << thread_id << ": Failed to write record batch: " 
                     << write_status.ToString() << std::endl;
            return false;
        }
        
        return true;
    } catch (const std::exception& e) {
        std::cerr << "Thread " << thread_id << ": Exception processing Arrow data: " 
                 << e.what() << std::endl;
        return false;
    }
}

// Process Arrow chunks to disk
void process_arrow_chunk_to_disk(
    const char* data,
    const std::vector<std::pair<size_t, size_t>>& indices,
    size_t start_idx,
    size_t end_idx,
    std::atomic<size_t>& progress_counter,
    simdjson::ondemand::parser& parser,
    std::mutex& output_mutex,
    std::string& last_json,
    size_t thread_id,
    const std::string& output_dir,
    const std::string& base_filename,
    size_t total_threads,
    std::shared_ptr<arrow::Schema> schema = nullptr
) {
    // Thread-specific output file
    std::string arrow_file_path = fs::path(output_dir) / fs::path(base_filename + "_" + std::to_string(thread_id) + ".arrow");
    
    // Create output file
    std::shared_ptr<arrow::io::FileOutputStream> output;
    auto output_result = arrow::io::FileOutputStream::Open(arrow_file_path);
    if (!output_result.ok()) {
        std::cerr << "Thread " << thread_id << ": Failed to open output file: " 
                 << output_result.status().ToString() << std::endl;
        return;
    }
    output = output_result.ValueOrDie();
    
    // Create Arrow writer - will be initialized with the first batch
    std::shared_ptr<arrow::ipc::RecordBatchWriter> writer;
    
    // Optimize buffer size based on data properties
    std::string jsonl_buffer;
    jsonl_buffer.reserve(16 * 1024 * 1024); // 16MB initial buffer
    
    // Adaptive batch sizing
    size_t records_to_process = end_idx - start_idx;
    size_t target_batch_size = std::min(size_t(50000), std::max(size_t(5000), records_to_process / 20));
    size_t max_batch_bytes = 64 * 1024 * 1024; // 64MB
    
    size_t current_batch_rows = 0;
    size_t current_batch_bytes = 0;
    size_t total_valid_rows = 0;
    size_t total_invalid_rows = 0;
    
    // Create memory pool
    arrow::MemoryPool* pool = arrow::default_memory_pool();
    
    // Process lines in batches with progress tracking
    size_t progress_batch_size = 0;
    const size_t PROGRESS_UPDATE_THRESHOLD = 10000;
    
    // Pre-scanning to estimate optimal batch size
    size_t sample_size = 0;
    size_t sample_bytes = 0;
    const size_t MAX_SAMPLES = 100;
    
    for (size_t i = start_idx; i < std::min(start_idx + MAX_SAMPLES, end_idx); i++) {
        const auto& [line_start, line_end] = indices[i];
        sample_size++;
        sample_bytes += (line_end - line_start);
    }
    
    // If we have samples, use them to estimate better batch sizes
    if (sample_size > 0) {
        size_t avg_line_size = sample_bytes / sample_size;
        target_batch_size = std::min(max_batch_bytes / avg_line_size, size_t(100000));
        target_batch_size = std::max(target_batch_size, size_t(1000));
    }
    
    // Process lines
    for (size_t i = start_idx; i < end_idx; ++i) {
        const auto& [line_start, line_end] = indices[i];
        const size_t line_length = line_end - line_start;

        if (line_length > 0) {
            bool valid_json = false;
            
            try {
                // Use simdjson to parse the line to verify it's valid
                simdjson::padded_string padded_line(data + line_start, line_length);
                auto doc = parser.iterate(padded_line);
                valid_json = true;
                
                // Add to buffer - use direct append for efficiency
                jsonl_buffer.append(data + line_start, line_length);
                jsonl_buffer.push_back('\n');
                
                current_batch_rows++;
                current_batch_bytes += line_length + 1;
                total_valid_rows++;
                progress_batch_size++;
                
                // Save the last JSON if this is the last record
                if (i == end_idx - 1 && thread_id == total_threads - 1) {
                    std::lock_guard<std::mutex> lock(output_mutex);
                    last_json = std::string(data + line_start, line_length);
                }

            } catch (const simdjson::simdjson_error& e) {
                total_invalid_rows++;
                
                // Log invalid JSON only if it's a small percentage
                if (total_invalid_rows < 10 || (total_valid_rows > 0 && 
                    total_invalid_rows * 100 / (total_valid_rows + total_invalid_rows) < 5)) {
                    std::cerr << "Thread " << thread_id << ": Invalid JSON at row " << (i - start_idx) 
                             << " - " << e.what() << std::endl;
                }
            } catch (const std::exception& e) {
                // Log standard exceptions
                total_invalid_rows++;
            } catch (...) {
                // Catch unknown exceptions
                total_invalid_rows++;
            }
                    
            // Process batch when threshold reached
            bool should_process_batch = false;
            
            // More aggressive batching based on records to process
            if (records_to_process > 1000000) {
                // For very large datasets, use larger batches
                should_process_batch = current_batch_rows >= target_batch_size * 2 || 
                                      current_batch_bytes >= max_batch_bytes;
            } else {
                // Standard batching criteria
                should_process_batch = current_batch_rows >= target_batch_size || 
                                      current_batch_bytes >= max_batch_bytes;
            }
            
            // Always process the final batch
            if (i == end_idx - 1) {
                should_process_batch = true;
            }
            
            if (should_process_batch && current_batch_rows > 0) {
                // Process the batch and write to Arrow file
                if (!process_and_write_batch(jsonl_buffer, writer, output, pool, schema, thread_id)) {
                    std::cerr << "Thread " << thread_id << ": Failed to process batch" << std::endl;
                }
                
                // Update progress counter in larger batches
                if (progress_batch_size >= PROGRESS_UPDATE_THRESHOLD) {
                    progress_counter.fetch_add(progress_batch_size, std::memory_order_relaxed);
                    progress_batch_size = 0;
                }
                
                // Reset batch tracking
                jsonl_buffer.clear();
                jsonl_buffer.reserve(16 * 1024 * 1024);
                current_batch_rows = 0;
                current_batch_bytes = 0;
            }
        }
    }
    
    // Update progress for any remaining records
    if (progress_batch_size > 0) {
        progress_counter.fetch_add(progress_batch_size, std::memory_order_relaxed);
    }
    
    // Close the writer
    if (writer) {
        auto status = writer->Close();
        if (!status.ok()) {
            std::cerr << "Thread " << thread_id << ": Error closing Arrow writer: " 
                     << status.ToString() << std::endl;
        }
    }
    
    // Report stats
    if (total_invalid_rows > 0) {
        std::cerr << "Thread " << thread_id << ": Processed " << (total_valid_rows + total_invalid_rows) 
                 << " lines, " << total_invalid_rows << " invalid JSON objects skipped (" 
                 << (total_invalid_rows * 100.0 / (total_valid_rows + total_invalid_rows)) << "%)" << std::endl;
    }
}

// Process Arrow chunks in memory
void process_arrow_chunk_in_memory(
    const char* data,
    const std::vector<std::pair<size_t, size_t>>& indices,
    size_t start_idx,
    size_t end_idx,
    std::atomic<size_t>& progress_counter,
    simdjson::ondemand::parser& parser,
    std::mutex& output_mutex,
    std::string& last_json,
    size_t thread_id,
    std::vector<ArrowMemoryChunk>& memory_chunks,
    size_t total_threads,
    std::shared_ptr<arrow::Schema> schema = nullptr
) {
    // Thread-specific memory chunk
    ArrowMemoryChunk memory_chunk;
    memory_chunk.schema = schema;
    
    // Optimize buffer size based on data properties
    std::string jsonl_buffer;
    jsonl_buffer.reserve(16 * 1024 * 1024); // 16MB initial buffer
    if (thread_id == total_threads - 1) {
        jsonl_buffer.reserve(16 * 1024 * 1024 * 10); // 160MB for the last thread
    }
    
    // Adaptive batch sizing
    size_t records_to_process = end_idx - start_idx;
    size_t target_batch_size = std::min(size_t(50000), std::max(size_t(5000), records_to_process / 20));
    size_t max_batch_bytes = 64 * 1024 * 1024; // 64MB
    
    size_t current_batch_rows = 0;
    size_t current_batch_bytes = 0;
    size_t total_valid_rows = 0;
    size_t total_invalid_rows = 0;
    
    // Create memory pool
    arrow::MemoryPool* pool = arrow::default_memory_pool();
    
    // Process lines in batches with progress tracking
    size_t progress_batch_size = 0;
    const size_t PROGRESS_UPDATE_THRESHOLD = 10000;
    
    // Pre-scanning to estimate optimal batch size
    size_t sample_size = 0;
    size_t sample_bytes = 0;
    const size_t MAX_SAMPLES = 100;
    
    for (size_t i = start_idx; i < std::min(start_idx + MAX_SAMPLES, end_idx); i++) {
        const auto& [line_start, line_end] = indices[i];
        sample_size++;
        sample_bytes += (line_end - line_start);
    }
    
    // If we have samples, use them to estimate better batch sizes
    if (sample_size > 0) {
        size_t avg_line_size = sample_bytes / sample_size;
        target_batch_size = std::min(max_batch_bytes / avg_line_size, size_t(100000));
        target_batch_size = std::max(target_batch_size, size_t(1000));
    }
    
    // Process lines
    for (size_t i = start_idx; i < end_idx; ++i) {
        const auto& [line_start, line_end] = indices[i];
        const size_t line_length = line_end - line_start;

        if (line_length > 0) {
            bool valid_json = false;
            
            try {
                // Use simdjson to parse the line to verify it's valid
                simdjson::padded_string padded_line(data + line_start, line_length);
                auto doc = parser.iterate(padded_line);
                valid_json = true;
                
                // Add to buffer - use direct append for efficiency
                jsonl_buffer.append(data + line_start, line_length);
                jsonl_buffer.push_back('\n');
                
                current_batch_rows++;
                current_batch_bytes += line_length + 1;
                total_valid_rows++;
                progress_batch_size++;
                
                // Save the last JSON if this is the last record
                if (i == end_idx - 1 && thread_id == total_threads - 1) {
                    std::lock_guard<std::mutex> lock(output_mutex);
                    last_json = std::string(data + line_start, line_length);
                }

            } catch (const simdjson::simdjson_error& e) {
                total_invalid_rows++;
                
                // Log invalid JSON only if it's a small percentage
                if (total_invalid_rows < 10 || (total_valid_rows > 0 && 
                    total_invalid_rows * 100 / (total_valid_rows + total_invalid_rows) < 5)) {
                    std::cerr << "Thread " << thread_id << ": Invalid JSON at row " << (i - start_idx) 
                             << " - " << e.what() << std::endl;
                }
            } catch (const std::exception& e) {
                // Log standard exceptions with their messages
                total_invalid_rows++;
                
                if (total_invalid_rows < 20) {
                    std::cerr << "Thread " << thread_id << ": Exception at row " << (i - start_idx) 
                            << " - " << e.what() << std::endl;
                }
            } catch (...) {
                // Still catch unknown exceptions but log them better
                total_invalid_rows++;
                
                if (total_invalid_rows < 20) {
                    std::cerr << "Thread " << thread_id << ": Unknown exception type at row " << (i - start_idx) << std::endl;
                }
            }
                    
            // Process batch when threshold reached
            bool should_process_batch = false;
            
            // More aggressive batching based on records to process
            if (records_to_process > 1000000) {
                // For very large datasets, use larger batches
                should_process_batch = current_batch_rows >= target_batch_size * 2 || 
                                      current_batch_bytes >= max_batch_bytes;
            } else {
                // Standard batching criteria
                should_process_batch = current_batch_rows >= target_batch_size || 
                                      current_batch_bytes >= max_batch_bytes;
            }
            
            // Always process the final batch
            if (i == end_idx - 1) {
                should_process_batch = true;
            }
            
            if (should_process_batch && current_batch_rows > 0) {
                // Convert the buffer to an Arrow batch and store in memory
                try {
                    // Create an input stream from the buffer with zero-copy where possible
                    auto buffer_ptr = std::make_shared<arrow::Buffer>(
                        reinterpret_cast<const uint8_t*>(jsonl_buffer.data()), 
                        jsonl_buffer.size()
                    );
                    auto input = std::make_shared<arrow::io::BufferReader>(buffer_ptr);
                    
                    // Configure JSON reader options - disable internal threading to avoid contention
                    auto read_options = arrow::json::ReadOptions::Defaults();
                    read_options.use_threads = false;
                    read_options.block_size = std::min(static_cast<int64_t>(jsonl_buffer.size()), 
                                                     static_cast<int64_t>(8 * 1024 * 1024));
                    
                    // Set parse options with schema if available
                    auto parse_options = arrow::json::ParseOptions::Defaults();
                    if (schema) {
                        parse_options.unexpected_field_behavior = arrow::json::UnexpectedFieldBehavior::Ignore;
                    }
                    
                    // Create JSON reader
                    auto reader_result = arrow::json::TableReader::Make(pool, input, read_options, parse_options);
                    if (!reader_result.ok()) {
                        std::cerr << "Thread " << thread_id << ": Failed to create JSON reader: " 
                                 << reader_result.status().ToString() << std::endl;
                    } else {
                        auto reader = reader_result.ValueOrDie();
                        
                        // Read the table
                        auto table_result = reader->Read();
                        if (!table_result.ok()) {
                            std::cerr << "Thread " << thread_id << ": Failed to read JSON data: " 
                                     << table_result.status().ToString() << std::endl;
                        } else {
                            auto table = table_result.ValueOrDie();
                            
                            // Save schema if we don't have one yet
                            if (!memory_chunk.schema) {
                                memory_chunk.schema = table->schema();
                            }
                            
                            // Convert table to batch
                            auto batch_result = table->CombineChunksToBatch();
                            if (!batch_result.ok()) {
                                std::cerr << "Thread " << thread_id << ": Failed to create record batch: " 
                                         << batch_result.status().ToString() << std::endl;
                            } else {
                                auto batch = batch_result.ValueOrDie();
                                
                                // Store the batch in memory
                                memory_chunk.batches.push_back(batch);
                            }
                        }
                    }
                } catch (const std::exception& e) {
                    std::cerr << "Thread " << thread_id << ": Exception processing Arrow data: " 
                             << e.what() << std::endl;
                }
                
                // Update progress counter in larger batches
                if (progress_batch_size >= PROGRESS_UPDATE_THRESHOLD) {
                    progress_counter.fetch_add(progress_batch_size, std::memory_order_relaxed);
                    progress_batch_size = 0;
                }
                
                // Reset batch tracking
                jsonl_buffer.clear();
                jsonl_buffer.reserve(16 * 1024 * 1024);
                current_batch_rows = 0;
                current_batch_bytes = 0;
            }
        }
    }
    
    // Update progress for any remaining records
    if (progress_batch_size > 0) {
        progress_counter.fetch_add(progress_batch_size, std::memory_order_relaxed);
    }
    
    // Store this thread's memory chunk in the shared vector
    {
        std::lock_guard<std::mutex> lock(output_mutex);
        memory_chunks[thread_id] = std::move(memory_chunk);
    }
    
    // Report stats
    if (total_invalid_rows > 0) {
        std::cerr << "Thread " << thread_id << ": Processed " << (total_valid_rows + total_invalid_rows) 
                 << " lines, " << total_invalid_rows << " invalid JSON objects skipped (" 
                 << (total_invalid_rows * 100.0 / (total_valid_rows + total_invalid_rows)) << "%)" << std::endl;
    }
}

// Function to merge in-memory Arrow chunks directly
arrow::Status merge_arrow_chunks_in_memory(
    const std::vector<ArrowMemoryChunk>& memory_chunks,
    const std::string& output_file,
    bool show_progress = true
) {
    if (memory_chunks.empty()) {
        return arrow::Status::OK();
    }
    
    // Start timing
    auto start_time = std::chrono::high_resolution_clock::now();
    
    std::cout << "Merging Arrow chunks in memory" << std::endl;
    // Find the first valid schema
    std::shared_ptr<arrow::Schema> schema;
    for (const auto& chunk : memory_chunks) {
        if (chunk.schema) {
            schema = chunk.schema;
            break;
        }
    }
    
    if (!schema) {
        return arrow::Status::Invalid("No valid schema found in memory chunks");
    }

    std::cout << "Schema found: " << schema->ToString() << std::endl;
    
    // Calculate total batches for progress tracking
    int64_t total_batches = 0;
    for (const auto& chunk : memory_chunks) {
        total_batches += chunk.batches.size();
    }
    
    std::cout << "Total batches: " << total_batches << std::endl;
    
    // Open output file
    ARROW_ASSIGN_OR_RAISE(auto output, arrow::io::FileOutputStream::Open(output_file));
    
    // Create the writer
    auto options = arrow::ipc::IpcWriteOptions::Defaults();
    options.write_legacy_ipc_format = false;
    options.max_recursion_depth = 128;
    
    ARROW_ASSIGN_OR_RAISE(auto writer, arrow::ipc::MakeFileWriter(output.get(), schema, options));
    
    // Progress tracking
    std::atomic<int64_t> processed_batches(0);
    std::atomic<bool> done(false);

    std::cout << "Starting progress reporting thread" << std::endl;
    // Start progress reporting thread if requested
    std::unique_ptr<std::thread> progress_thread;
    if (show_progress) {
        progress_thread = std::make_unique<std::thread>([&]() {
            auto last_update_time = std::chrono::high_resolution_clock::now();
            int64_t last_batches = 0;
            
            while (!done.load()) {
                std::this_thread::sleep_for(std::chrono::milliseconds(250));
                
                auto current_time = std::chrono::high_resolution_clock::now();
                auto elapsed_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
                    current_time - last_update_time).count();
                
                int64_t current_batches = processed_batches.load();
                
                // Calculate batches per second
                double batches_per_sec = 0;
                if (elapsed_ms > 0) {
                    batches_per_sec = (current_batches - last_batches) * 1000.0 / elapsed_ms;
                }
                
                // Calculate overall progress percentage
                double progress_pct = total_batches > 0 ? 
                    (current_batches * 100.0 / total_batches) : 0;
                
                // Calculate ETA
                std::string eta_str = "calculating...";
                if (batches_per_sec > 0 && current_batches < total_batches) {
                    double eta_seconds = (total_batches - current_batches) / batches_per_sec;
                    if (eta_seconds > 0) {
                        if (eta_seconds < 60) {
                            eta_str = std::to_string(static_cast<int>(eta_seconds)) + "s";
                        } else {
                            int mins = static_cast<int>(eta_seconds / 60);
                            int secs = static_cast<int>(eta_seconds) % 60;
                            eta_str = std::to_string(mins) + "m " + std::to_string(secs) + "s";
                        }
                    }
                }
                
                std::cout << "\rMerging Arrow chunks in memory: " << std::fixed << std::setprecision(1) 
                         << progress_pct << "% - " 
                         << current_batches << "/" << total_batches << " batches | "
                         << std::setprecision(1) << batches_per_sec << " batches/sec | ETA: " 
                         << eta_str << "        " << std::flush;
                
                last_update_time = current_time;
                last_batches = current_batches;
            }
            std::cout << std::endl;
        });
    }
    
    // Process all chunks in order
    for (size_t chunk_idx = 0; chunk_idx < memory_chunks.size(); chunk_idx++) {
        const auto& chunk = memory_chunks[chunk_idx];
        
        // Process all batches in this chunk
        for (const auto& batch : chunk.batches) {
            // Write the batch
            auto status = writer->WriteRecordBatch(*batch);
            if (!status.ok()) {
                std::cerr << "Error writing batch from chunk " << chunk_idx << ": " 
                         << status.ToString() << std::endl;
                continue;
            }
            
            processed_batches.fetch_add(1);
        }
    }
    
    // Signal progress thread to complete
    done.store(true);
    if (progress_thread && progress_thread->joinable()) {
        progress_thread->join();
    }
    
    // Close the writer
    ARROW_RETURN_NOT_OK(writer->Close());
    
    // Report timing
    auto end_time = std::chrono::high_resolution_clock::now();
    auto duration_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
        end_time - start_time).count();
    
    if (show_progress) {
        double batches_per_sec = total_batches / (duration_ms / 1000.0 + 0.001); // Avoid division by zero
        
        std::cout << "Merged " << memory_chunks.size() << " in-memory chunks with " 
                 << total_batches << " batches in " 
                 << duration_ms / 1000.0 << " seconds (" 
                 << std::fixed << std::setprecision(2) << batches_per_sec << " batches/s)" << std::endl;
    }
    
    return arrow::Status::OK();
}

// Function to merge Arrow files from disk
arrow::Status merge_arrow_files(
    const std::vector<std::string>& file_paths,
    const std::string& output_file,
    bool show_progress = true
) {
    if (file_paths.empty()) {
        return arrow::Status::OK();
    }
    
    // Start timing
    auto start_time = std::chrono::high_resolution_clock::now();
    
    // Open the first file to get schema
    ARROW_ASSIGN_OR_RAISE(auto input_stream, arrow::io::ReadableFile::Open(file_paths[0]));
    ARROW_ASSIGN_OR_RAISE(auto reader, arrow::ipc::RecordBatchFileReader::Open(input_stream));
    auto schema = reader->schema();
    
    // Open output file
    ARROW_ASSIGN_OR_RAISE(auto output, arrow::io::FileOutputStream::Open(output_file));
    
    // Create the writer with optimized options
    auto write_options = arrow::ipc::IpcWriteOptions::Defaults();
    write_options.write_legacy_ipc_format = false;
    write_options.max_recursion_depth = 128;
    
    ARROW_ASSIGN_OR_RAISE(auto writer, arrow::ipc::MakeFileWriter(output.get(), schema, write_options));
    
    // Progress tracking
    int64_t total_batches = 0;
    std::atomic<int64_t> processed_batches(0);
    std::atomic<bool> done(false);
    
    // First pass to count total batches for progress reporting
    if (show_progress) {
        for (const auto& file_path : file_paths) {
            ARROW_ASSIGN_OR_RAISE(auto file, arrow::io::ReadableFile::Open(file_path));
            ARROW_ASSIGN_OR_RAISE(auto file_reader, arrow::ipc::RecordBatchFileReader::Open(file));
            total_batches += file_reader->num_record_batches();
        }
    }
    
    // Start progress reporting thread if requested
    std::unique_ptr<std::thread> progress_thread;
    if (show_progress) {
        progress_thread = std::make_unique<std::thread>([&]() {
            auto last_update_time = std::chrono::high_resolution_clock::now();
            int64_t last_batches = 0;
            
            while (!done.load()) {
                std::this_thread::sleep_for(std::chrono::milliseconds(250));
                
                auto current_time = std::chrono::high_resolution_clock::now();
                auto elapsed_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
                    current_time - last_update_time).count();
                
                int64_t current_batches = processed_batches.load();
                
                // Calculate batches per second
                double batches_per_sec = 0;
                if (elapsed_ms > 0) {
                    batches_per_sec = (current_batches - last_batches) * 1000.0 / elapsed_ms;
                }
                
                // Calculate overall progress percentage
                double progress_pct = total_batches > 0 ? 
                    (current_batches * 100.0 / total_batches) : 0;
                
                // Calculate ETA
                std::string eta_str = "calculating...";
                if (batches_per_sec > 0 && current_batches < total_batches) {
                    double eta_seconds = (total_batches - current_batches) / batches_per_sec;
                    if (eta_seconds > 0) {
                        if (eta_seconds < 60) {
                            eta_str = std::to_string(static_cast<int>(eta_seconds)) + "s";
                        } else {
                            int mins = static_cast<int>(eta_seconds / 60);
                            int secs = static_cast<int>(eta_seconds) % 60;
                            eta_str = std::to_string(mins) + "m " + std::to_string(secs) + "s";
                        }
                    }
                }
                
                std::cout << "\rMerging Arrow files: " << std::fixed << std::setprecision(1) 
                         << progress_pct << "% - " 
                         << current_batches << "/" << total_batches << " batches | "
                         << std::setprecision(1) << batches_per_sec << " batches/sec | ETA: " 
                         << eta_str << "        " << std::flush;
                
                last_update_time = current_time;
                last_batches = current_batches;
            }
            std::cout << std::endl;
        });
    }
    
    // Process all files
    for (const auto& file_path : file_paths) {
        // Open input file
        ARROW_ASSIGN_OR_RAISE(auto file, arrow::io::ReadableFile::Open(file_path));
        ARROW_ASSIGN_OR_RAISE(auto file_reader, arrow::ipc::RecordBatchFileReader::Open(file));
        
        // Read and write each batch
        for (int i = 0; i < file_reader->num_record_batches(); i++) {
            ARROW_ASSIGN_OR_RAISE(auto batch, file_reader->ReadRecordBatch(i));
            ARROW_RETURN_NOT_OK(writer->WriteRecordBatch(*batch));
            processed_batches.fetch_add(1);
        }
    }
    
    // Signal progress thread to complete
    done.store(true);
    if (progress_thread && progress_thread->joinable()) {
        progress_thread->join();
    }
    
    // Close the writer
    ARROW_RETURN_NOT_OK(writer->Close());
    
    // Report timing
    auto end_time = std::chrono::high_resolution_clock::now();
    auto duration_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
        end_time - start_time).count();
    
    if (show_progress) {
        double batches_per_sec = processed_batches.load() / (duration_ms / 1000.0 + 0.001); // Avoid division by zero
        
        std::cout << "Merged " << file_paths.size() << " Arrow files with " 
                 << processed_batches.load() << " batches in " 
                 << duration_ms / 1000.0 << " seconds (" 
                 << std::fixed << std::setprecision(2) << batches_per_sec << " batches/s)" << std::endl;
    }
    
    return arrow::Status::OK();
}

// Merge JSONL files using direct I/O
bool merge_jsonl_files(
    const std::vector<std::string>& file_paths, 
    const std::string& output_file,
    bool show_progress = true
) {
    // Start timing
    auto start_time = std::chrono::high_resolution_clock::now();
    
    // Use direct I/O with a large buffer for merging
    int out_fd = open(output_file.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644);
    if (out_fd == -1) {

        std::cerr << "Failed to open output file for merging: " << output_file 
                 << " - " << strerror(errno) << std::endl;
        return false;
    }
    
    // Setup progress tracking
    size_t total_bytes = 0;
    if (show_progress) {
        for (const auto& file : file_paths) {
            struct stat sb;
            if (stat(file.c_str(), &sb) == 0) {
                total_bytes += sb.st_size;
            }
        }
    }
    
    std::atomic<size_t> processed_bytes(0);
    std::atomic<bool> done(false);
    
    // Progress reporting thread
    std::unique_ptr<std::thread> progress_thread;
    if (show_progress) {
        progress_thread = std::make_unique<std::thread>([&]() {
            auto last_update_time = std::chrono::high_resolution_clock::now();
            size_t last_bytes = 0;
            
            while (!done.load()) {
                std::this_thread::sleep_for(std::chrono::milliseconds(250));
                
                auto current_time = std::chrono::high_resolution_clock::now();
                auto elapsed_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
                    current_time - last_update_time).count();
                
                size_t current_bytes = processed_bytes.load();
                
                // Calculate MB per second
                double mb_per_sec = 0;
                if (elapsed_ms > 0) {
                    mb_per_sec = ((current_bytes - last_bytes) / (1024.0 * 1024.0)) * (1000.0 / elapsed_ms);
                }
                
                // Calculate overall progress percentage
                double progress_pct = total_bytes > 0 ? 
                    (current_bytes * 100.0 / total_bytes) : 0;
                
                // Calculate ETA
                std::string eta_str = "calculating...";
                if (mb_per_sec > 0 && current_bytes < total_bytes) {
                    double remaining_mb = (total_bytes - current_bytes) / (1024.0 * 1024.0);
                    double eta_seconds = remaining_mb / mb_per_sec;
                    if (eta_seconds > 0) {
                        if (eta_seconds < 60) {
                            eta_str = std::to_string(static_cast<int>(eta_seconds)) + "s";
                        } else {
                            int mins = static_cast<int>(eta_seconds / 60);
                            int secs = static_cast<int>(eta_seconds) % 60;
                            eta_str = std::to_string(mins) + "m " + std::to_string(secs) + "s";
                        }
                    }
                }
                
                std::cout << "\rMerging JSONL files: " << std::fixed << std::setprecision(1) 
                         << progress_pct << "% - " 
                         << (current_bytes / (1024 * 1024)) << "/" << (total_bytes / (1024 * 1024)) << " MB | "
                         << std::setprecision(2) << mb_per_sec << " MB/sec | ETA: " 
                         << eta_str << "        " << std::flush;
                
                last_update_time = current_time;
                last_bytes = current_bytes;
            }
            std::cout << std::endl;
        });
    }
    
    constexpr size_t MERGE_BUFFER_SIZE = LARGE_BUFFER_SIZE; // 256MB buffer
    std::unique_ptr<char[]> buffer(new char[MERGE_BUFFER_SIZE]);
    
    for (const auto& file_path : file_paths) {
        int in_fd = open(file_path.c_str(), O_RDONLY);
        if (in_fd == -1) {
            std::cerr << "Failed to open chunk file: " << file_path 
                     << " - " << strerror(errno) << std::endl;
            close(out_fd);
            
            // Signal progress thread to complete
            done.store(true);
            if (progress_thread && progress_thread->joinable()) {
                progress_thread->join();
            }
            
            return false;
        }
        
        // Copy file contents with large buffer
        ssize_t bytes_read;
        while ((bytes_read = read(in_fd, buffer.get(), MERGE_BUFFER_SIZE)) > 0) {
            write(out_fd, buffer.get(), bytes_read);
            processed_bytes.fetch_add(bytes_read);
        }
        
        close(in_fd);
    }
    
    close(out_fd);
    
    // Signal progress thread to complete
    done.store(true);
    if (progress_thread && progress_thread->joinable()) {
        progress_thread->join();
    }
    
    // Report timing
    auto end_time = std::chrono::high_resolution_clock::now();
    auto duration_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
        end_time - start_time).count();
    
    if (show_progress) {
        double mb_per_sec = (processed_bytes.load() / (1024.0 * 1024.0)) / (duration_ms / 1000.0);
        
        std::cout << "Merged " << file_paths.size() << " JSONL files, " 
                 << (processed_bytes.load() / (1024 * 1024)) << " MB in "
                 << duration_ms / 1000.0 << " seconds (" 
                 << std::fixed << std::setprecision(2) << mb_per_sec << " MB/s)" << std::endl;
    }
    
    return true;
}


int main(int argc, char* argv[]) {
    try {
        // Default parameters
        std::string file_path = "/tmp/test_data.jsonl";
        std::string output_dir = "/mnt/localdisk/parser-output";
        std::string base_filename = "parsed_chunk";
        bool merge_files = true;
        auto now = std::chrono::system_clock::now();
        auto time_t_now = std::chrono::system_clock::to_time_t(now);
        std::stringstream timestamp;
        timestamp << std::put_time(std::localtime(&time_t_now), "%Y-%m-%d_%H-%M");
        std::string output_file = "parsed_results_" + timestamp.str() + ".arrow";

        bool use_arrow = true;  // Flag for Arrow format output
        ProcessingMode processing_mode = ProcessingMode::Auto; // Auto-detect memory vs. disk
        bool cleanup_files = true;  // Cleanup temporary files by default
        bool quiet = false;  // Show progress by default
        
        // Parse command line arguments
        for (int i = 1; i < argc; i++) {
            std::string arg = argv[i];
            if (arg == "--input" && i + 1 < argc) {
                file_path = argv[++i];
            } else if (arg == "--output-dir" && i + 1 < argc) {
                output_dir = argv[++i];
            } else if (arg == "--base-filename" && i + 1 < argc) {
                base_filename = argv[++i];
            } else if (arg == "--output" && i + 1 < argc) {
                output_file = argv[++i];
            } else if (arg == "--merge") {
                merge_files = true;
            } else if (arg == "--arrow") {
                use_arrow = true;
                // Adjust output file extension if needed
                if (output_file.find(".jsonl") != std::string::npos && output_file.rfind(".jsonl") == output_file.length() - 6) {
                    output_file = output_file.substr(0, output_file.length() - 6) + ".arrow";
                } else if (output_file.find(".json") != std::string::npos && output_file.rfind(".json") == output_file.length() - 5) {
                    output_file = output_file.substr(0, output_file.length() - 5) + ".arrow";
                } else {
                    output_file += ".arrow";
                }
            } else if (arg == "--in-memory") {
                processing_mode = ProcessingMode::ForceMemory;
            } else if (arg == "--disk") {
                processing_mode = ProcessingMode::ForceDisk;
            } else if (arg == "--no-cleanup") {
                cleanup_files = false;
            } else if (arg == "--quiet") {
                quiet = true;
            } else if (arg == "--help") {
                std::cout << "JSONL Parser with Memory Optimization v2.0" << std::endl;
                std::cout << "Usage: " << argv[0] << " [options]" << std::endl;
                std::cout << "Options:" << std::endl;
                std::cout << "  --input FILE          Input JSONL file (default: /tmp/test_data.jsonl)" << std::endl;
                std::cout << "  --output-dir DIR      Directory to save output files (default: ./parser-output)" << std::endl;
                std::cout << "  --base-filename NAME  Base name for output chunk files (default: parsed_chunk)" << std::endl;
                std::cout << "  --output FILE         Final output file for merged results (default: parsed_results.jsonl)" << std::endl;
                std::cout << "  --merge               Merge output chunks into a single file (default: false)" << std::endl;
                std::cout << "  --arrow               Output in Apache Arrow format instead of JSONL (default: false)" << std::endl;
                std::cout << "  --in-memory           Force in-memory processing (default: auto-detect)" << std::endl;
                std::cout << "  --disk                Force disk-based processing (default: auto-detect)" << std::endl;
                std::cout << "  --no-cleanup          Don't remove temporary files after merging (default: cleanup)" << std::endl;
                std::cout << "  --quiet               Suppress progress information (default: show progress)" << std::endl;
                std::cout << "  --help                Show this help message" << std::endl;
                std::cout << std::endl;
                std::cout << "Memory Management:" << std::endl;
                std::cout << "  By default, the program will automatically choose between in-memory" << std::endl;
                std::cout << "  and disk-based processing based on file size and available system memory." << std::endl;
                std::cout << "  - Small files (<256MB): Always use in-memory processing" << std::endl;
                std::cout << "  - Medium files: Use in-memory if enough RAM is available" << std::endl;
                std::cout << "  - Large files (>2GB): Use disk-based processing unless plenty of RAM" << std::endl;
                return 0;
            }
        }
        
        // Create output directory if it doesn't exist
        fs::create_directories(output_dir);
        
        auto start_time = std::chrono::high_resolution_clock::now();

        // Memory map the file
        MappedFile file(file_path);
        
        if (!quiet) {
            std::cout << "File size: " << file.size / (1024 * 1024) << " MB" << std::endl;
        }

        // Determine optimal processing mode based on file size and memory
        processing_mode = determine_processing_mode(file.size, processing_mode);
        bool in_memory_processing = (processing_mode == ProcessingMode::ForceMemory);
        
        if (!quiet) {
            if (in_memory_processing) {
                std::cout << "Using in-memory processing" << std::endl;
            } else {
                std::cout << "Using disk-based processing" << std::endl;
            }
        }

        // Determine optimal thread count based on file size and system
        const size_t num_threads = calculate_optimal_threads(file.size);
        if (!quiet) {
            std::cout << "Using " << num_threads << " threads" << std::endl;
        }

        // Find all line indices with parallel scan
        if (!quiet) {
            std::cout << "Finding line indices using parallel scan..." << std::endl;
        }
        auto indices_start = std::chrono::high_resolution_clock::now();
        auto indices = find_line_indices_optimized(file, num_threads);
        auto indices_end = std::chrono::high_resolution_clock::now();
        auto indices_duration = std::chrono::duration_cast<std::chrono::milliseconds>(
            indices_end - indices_start).count() / 1000.0;
        
        const size_t total_lines = indices.size();
        if (!quiet) {
            std::cout << "Found " << total_lines << " lines in " << indices_duration << " seconds" << std::endl;
        }

        // Calculate lines per thread
        const size_t lines_per_thread = total_lines / num_threads;

        // Create thread pool
        std::vector<std::thread> threads;
        std::atomic<size_t> progress_counter(0);
        std::mutex output_mutex;
        std::string last_json;

        // Start progress reporting thread if not in quiet mode
        std::atomic<bool> progress_done(false);
        std::thread progress_thread;
        
        if (!quiet) {
            progress_thread = std::thread([&]() {
                size_t last_progress = 0;
                auto start = std::chrono::high_resolution_clock::now();
                
                // Allocate a vector for rate tracking
                const size_t rate_window = 10;  // Track last 10 progress updates
                std::vector<double> recent_rates;

                while (!progress_done.load(std::memory_order_relaxed)) {
                    std::this_thread::sleep_for(std::chrono::milliseconds(500));

                    size_t current = progress_counter.load(std::memory_order_relaxed);
                    auto now = std::chrono::high_resolution_clock::now();
                    auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(now - start).count();
                    auto total_elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(now - start_time).count();

                    if (current > last_progress && elapsed > 0) {
                        // Calculate recent rate
                        double rate = (current - last_progress) * 1000.0 / elapsed;
                        
                        // Keep track of recent rates for smoothing
                        recent_rates.push_back(rate);
                        if (recent_rates.size() > rate_window) {
                            recent_rates.erase(recent_rates.begin());
                        }
                        
                        // Calculate average rate
                        double avg_rate = 0;
                        for (double r : recent_rates) {
                            avg_rate += r;
                        }
                        avg_rate /= recent_rates.size();
                        
                        // Calculate overall rate
                        double overall_rate = current * 1000.0 / total_elapsed;
                        
                        // Calculate ETA
                        double eta_seconds = 0;
                        if (avg_rate > 0 && current < total_lines) {
                            eta_seconds = (total_lines - current) / avg_rate;
                        }
                        
                        // Format ETA string
                        std::string eta_str;
                        if (eta_seconds > 0) {
                            int eta_mins = static_cast<int>(eta_seconds) / 60;
                            int eta_secs = static_cast<int>(eta_seconds) % 60;
                            eta_str = " [ETA: " + std::to_string(eta_mins) + "m " + 
                                     std::to_string(eta_secs) + "s]";
                        }
                        
                        double percent = current * 100.0 / total_lines;

                        std::cout << "\rProgress: " << std::fixed << std::setprecision(1)
                                  << percent << "% - Processed " << current << " of "
                                  << total_lines << " records (" << static_cast<size_t>(avg_rate)
                                  << " records/sec)" << eta_str << "           " << std::flush;

                        last_progress = current;
                        start = now;
                    }
                }
                std::cout << std::endl;
            });
        }

        // Create parser for each thread
        std::vector<simdjson::ondemand::parser> parsers(num_threads);

        // Pre-infer schema for Arrow if needed
        std::shared_ptr<arrow::Schema> schema;
        if (use_arrow) {
            if (!quiet) {
                std::cout << "Inferring Arrow schema from sample data..." << std::endl;
            }
            auto schema_result = infer_schema_from_sample(
                static_cast<const char*>(file.data), indices);
            
            if (!schema_result.ok()) {
                std::cerr << "Failed to infer schema: " << schema_result.status().ToString() << std::endl;
            } else {
                schema = schema_result.ValueOrDie();
                if (!quiet) {
                    std::cout << "Schema inferred successfully, fields: " << schema->num_fields() << std::endl;
                }
            }
        }
        
        // Vector to hold memory chunks if using in-memory processing
        std::vector<ArrowMemoryChunk> memory_chunks;
        if (use_arrow && in_memory_processing) {
            memory_chunks.resize(num_threads);
        }

        // Select processing mode and launch threads
        if (use_arrow) {
            if (in_memory_processing) {
                if (!quiet) {
                    std::cout << "Processing in Arrow format and storing in memory" << std::endl;
                }
                
                // Start worker threads for in-memory Arrow processing
                for (size_t thread_id = 0; thread_id < num_threads; ++thread_id) {
                    const size_t start_idx = thread_id * lines_per_thread;
                    const size_t end_idx = (thread_id == num_threads - 1) ? total_lines : (thread_id + 1) * lines_per_thread;

                    threads.emplace_back(
                        process_arrow_chunk_in_memory,
                        static_cast<const char*>(file.data),
                        std::cref(indices),
                        start_idx,
                        end_idx,
                        std::ref(progress_counter),
                        std::ref(parsers[thread_id]),
                        std::ref(output_mutex),
                        std::ref(last_json),
                        thread_id,
                        std::ref(memory_chunks),
                        num_threads,
                        schema
                    );
                }
            } else {
                if (!quiet) {
                    std::cout << "Processing in Arrow format and writing to disk" << std::endl;
                }
                
                // Start worker threads for disk-based Arrow processing
                for (size_t thread_id = 0; thread_id < num_threads; ++thread_id) {
                    const size_t start_idx = thread_id * lines_per_thread;
                    const size_t end_idx = (thread_id == num_threads - 1) ? total_lines : (thread_id + 1) * lines_per_thread;

                    threads.emplace_back(
                        process_arrow_chunk_to_disk,
                        static_cast<const char*>(file.data),
                        std::cref(indices),
                        start_idx,
                        end_idx,
                        std::ref(progress_counter),
                        std::ref(parsers[thread_id]),
                        std::ref(output_mutex),
                        std::ref(last_json),
                        thread_id,
                        std::cref(output_dir),
                        std::cref(base_filename),
                        num_threads,
                        schema
                    );
                }
            }
        } else {
            if (!quiet) {
                std::cout << "Processing in JSONL format" << std::endl;
            }
            
            // Start worker threads for JSONL processing
            for (size_t thread_id = 0; thread_id < num_threads; ++thread_id) {
                const size_t start_idx = thread_id * lines_per_thread;
                const size_t end_idx = (thread_id == num_threads - 1) ? total_lines : (thread_id + 1) * lines_per_thread;

                threads.emplace_back(
                    process_jsonl_chunk,
                    static_cast<const char*>(file.data),
                    std::cref(indices),
                    start_idx,
                    end_idx,
                    std::ref(progress_counter),
                    std::ref(parsers[thread_id]),
                    std::ref(output_mutex),
                    std::ref(last_json),
                    thread_id,
                    std::cref(output_dir),
                    std::cref(base_filename),
                    num_threads
                );
            }
        }

        // Wait for worker threads to complete
        for (auto& thread : threads) {
            thread.join();
        }
        
        // Signal progress thread to complete
        progress_done.store(true, std::memory_order_relaxed);
        if (!quiet && progress_thread.joinable()) {
            progress_thread.join();
        }

        // Report results after processing
        auto processing_end_time = std::chrono::high_resolution_clock::now();
        auto processing_duration = std::chrono::duration_cast<std::chrono::milliseconds>(
            processing_end_time - start_time).count() / 1000.0;

        if (!quiet) {
            std::cout << "\nProcessed " << total_lines << " records in " << processing_duration << " seconds" << std::endl;
            std::cout << "Processing performance: " << (total_lines / processing_duration) << " records/second" << std::endl;
        }

        // Handle merging if requested
        if (merge_files) {
            std::string merged_file_path = fs::path(output_dir) / fs::path(output_file);
            
            if (use_arrow && in_memory_processing) {
                // Merge from in-memory Arrow chunks
                if (!quiet) {
                    std::cout << "Merging in-memory Arrow chunks..." << std::endl;
                }
                auto merge_start = std::chrono::high_resolution_clock::now();
                
                auto status = merge_arrow_chunks_in_memory(memory_chunks, merged_file_path, !quiet);
                
                auto merge_end = std::chrono::high_resolution_clock::now();
                auto merge_duration = std::chrono::duration_cast<std::chrono::milliseconds>(
                    merge_end - merge_start).count() / 1000.0;
                
                if (!status.ok()) {
                    std::cerr << "Error merging in-memory Arrow chunks: " << status.ToString() << std::endl;
                } else if (!quiet) {
                    std::cout << "Merged in-memory Arrow chunks in " << merge_duration << " seconds" << std::endl;
                    std::cout << "Final output saved to " << merged_file_path << std::endl;
                }
            } else if (use_arrow) {
                // Merge Arrow files from disk
                if (!quiet) {
                    std::cout << "Merging Arrow files from disk..." << std::endl;
                }
                auto merge_start = std::chrono::high_resolution_clock::now();
                
                // Collect all arrow chunk files
                std::vector<std::string> arrow_files;
                for (size_t i = 0; i < num_threads; i++) {
                    std::string file_path = fs::path(output_dir) / fs::path(base_filename + "_" + std::to_string(i) + ".arrow");
                    if (fs::exists(file_path)) {
                        arrow_files.push_back(file_path);
                    }
                }
                
                // Merge Arrow files
                auto status = merge_arrow_files(arrow_files, merged_file_path, !quiet);
                
                auto merge_end = std::chrono::high_resolution_clock::now();
                auto merge_duration = std::chrono::duration_cast<std::chrono::milliseconds>(
                    merge_end - merge_start).count() / 1000.0;
                
                if (!status.ok()) {
                    std::cerr << "Error merging Arrow files: " << status.ToString() << std::endl;
                } else {
                    // Remove chunk files if cleanup is enabled
                    if (cleanup_files) {
                        for (const auto& file : arrow_files) {
                            if (fs::exists(file)) {
                                fs::remove(file);
                            }
                        }
                    }
                    
                    if (!quiet) {
                        std::cout << "Merged Arrow files in " << merge_duration << " seconds" << std::endl;
                        std::cout << "Final output saved to " << merged_file_path << std::endl;
                    }
                }
            } else {
                // Merge JSONL files
                if (!quiet) {
                    std::cout << "Merging JSONL files..." << std::endl;
                }
                auto merge_start = std::chrono::high_resolution_clock::now();
                
                // Collect all JSONL chunk files
                std::vector<std::string> jsonl_files;
                for (size_t i = 0; i < num_threads; i++) {
                    std::string file_path = fs::path(output_dir) / fs::path(base_filename + "_" + std::to_string(i) + ".jsonl");
                    if (fs::exists(file_path)) {
                        jsonl_files.push_back(file_path);
                    }
                }
                
                // Merge JSONL files
                bool merge_success = merge_jsonl_files(jsonl_files, merged_file_path, !quiet);
                
                auto merge_end = std::chrono::high_resolution_clock::now();
                auto merge_duration = std::chrono::duration_cast<std::chrono::milliseconds>(
                    merge_end - merge_start).count() / 1000.0;
                
                if (!merge_success) {
                    std::cerr << "Error merging JSONL files" << std::endl;
                } else {
                    // Remove chunk files if cleanup is enabled
                    if (cleanup_files) {
                        for (const auto& file : jsonl_files) {
                            if (fs::exists(file)) {
                                fs::remove(file);
                            }
                        }
                    }
                    
                    if (!quiet) {
                        std::cout << "Merged JSONL files in " << merge_duration << " seconds" << std::endl;
                        std::cout << "Final output saved to " << merged_file_path << std::endl;
                    }
                }
            }
        } else if (!quiet) {
            // If no merging, report where the output is
            std::string extension = use_arrow ? ".arrow" : ".jsonl";
            std::cout << "Output saved as separate chunk files with " << extension 
                      << " extension in " << output_dir << std::endl;
        }

        // Total elapsed time
        auto end_time = std::chrono::high_resolution_clock::now();
        auto total_duration = std::chrono::duration_cast<std::chrono::milliseconds>(
            end_time - start_time).count() / 1000.0;

        if (!quiet) {
            std::cout << "Total execution time: " << total_duration << " seconds" << std::endl;

            // Print last record (truncated if too long)
            if (!last_json.empty()) {
                if (last_json.length() > 100) {
                    std::cout << "\nLast record: " << last_json.substr(0, 100) << "..." << std::endl;
                } else {
                    std::cout << "\nLast record: " << last_json << std::endl;
                }
            }
        }

    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return 1;
    } catch (...) {
        std::cerr << "Unknown error occurred" << std::endl;
        return 1;
    }

    return 0;
}