import modal
import json
import asyncio
import pandas as pd
import pickle
from tqdm.asyncio import tqdm
import numpy as np


async def batch_encode_with_concurrency_control(
    input_data_list: list, max_concurrent: int = 400
):
    f = modal.Function.lookup("ditto-v2-dev", "DittoWorkerStub.encode_audio")

    # Create semaphore to limit concurrency
    semaphore = asyncio.Semaphore(max_concurrent)

    async def process_single(input_json):
        async with semaphore:
            return await f.remote.aio(queue_item_json=input_json)

    # Process all items concurrently with limit and progress bar
    tasks = [process_single(input_json) for input_json in input_data_list]
    results = await tqdm.gather(*tasks, desc="Processing audio files")

    # Create result dictionary
    result_dict = {}
    for input_json, result in zip(input_data_list, results):
        parsed_input = json.loads(input_json)
        result_dict[parsed_input["id"]] = result

    return result_dict


PREF_PATH = "/app2/suno/data/sara/cover_pref_filtering/pariwise_preferences.pkl"
df = pd.read_pickle(PREF_PATH)

s3_ids = (
    np.column_stack([df["s3_id_1"], df["s3_id_2"], df["cover_clip_id"]])
    .flatten()
    .tolist()
)
print(f"Total S3 IDs: {len(s3_ids)}")

input_s3_ids = [json.dumps(dict(id=s3_id, metadata={})) for s3_id in s3_ids]

test_inputs = input_s3_ids  # [:10000]
print(f"Processing {len(test_inputs)} items...")

# Usage
result_dict = asyncio.run(batch_encode_with_concurrency_control(test_inputs))

print(f"Successfully processed {len(result_dict)} items")

pickle_path = "/app2/suno/data/sara/cover_pref_filtering/batch_encode_results.pkl"

with open(pickle_path, "wb") as f:
    pickle.dump(result_dict, f)

print(f"Results saved to {pickle_path}")