import json
import pyarrow as pa
import pyarrow.parquet as pq
import pandas as pd
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm


def process_item(item):
    # Read the image and convert it to byte format
    with open(item["image"], "rb") as img_file:
        img_bytes = img_file.read()

    record = {
        "image": img_bytes,
        "conversations": json.dumps(item["conversations"])  # Serialize as JSON string
    }
    return record


# Read the JSON file
with open('merged_half.json', 'r') as file:
    data = json.load(file)

local_path = 'merged_first_half.parquet'

# Get the number of CPU cores in the system
cpu_count = os.cpu_count()

# Process data in batches
batch_size = 100000  # Can be adjusted based on actual needs
num_batches = (len(data) + batch_size - 1) // batch_size

# Local file path
# local_path = 'final_data_4ch.parquet'

# Initialize ParquetWriter
with open(local_path, 'wb') as local_file:
    writer = None

    for batch_index in range(num_batches):
        start_index = batch_index * batch_size
        end_index = min((batch_index + 1) * batch_size, len(data))
        batch_data = data[start_index:end_index]

        # Use ThreadPoolExecutor for parallel processing
        records = []
        with ThreadPoolExecutor(max_workers=cpu_count) as executor:
            future_to_record = {executor.submit(process_item, item): item for item in batch_data}
            for future in tqdm(as_completed(future_to_record), total=len(future_to_record),
                               desc=f"Processing Batch {batch_index + 1}/{num_batches}"):
                try:
                    record = future.result()
                    records.append(record)
                except Exception as exc:
                    print(f'Generated an exception: {exc}')

        # Create a PyArrow table
        table = pa.Table.from_pandas(pd.DataFrame(records))

        # If it's the first batch, set the writer and schema
        if writer is None:
            writer = pq.ParquetWriter(local_file, table.schema, version='2.6', use_dictionary=True, compression='snappy')

        # Write to the Parquet file in chunks
        for i in tqdm(range(0, len(table), 4), desc=f"Writing Batch {batch_index + 1}/{num_batches} to Parquet"):
            writer.write_table(table.slice(i, 4))

    writer.close()

print("Completed: Batches saved as Parquet files to local directory")