import json
import logging
import math
from concurrent.futures import ThreadPoolExecutor, as_completed

import config
from oa_contract_ingest_to_abacus.connectors import s3
from oa_contract_ingest_to_abacus.utils import _get_min_contract_id, _get_max_contract_id

log = logging.getLogger(__name__)


def _upload_batch_to_s3(batch: list[tuple[int, dict]], batch_key: str, bucket) -> None:
    for account_to_json in batch:
        s3.create_object(
            bucket,
            f'{batch_key}/{account_to_json[0]}.json',
            json.dumps(account_to_json[1])
        )


def upload_file_in_batches(
        account_to_json: list[tuple[int, dict]],
        key_prefix: str,
        bucket,
) -> list[str]:
    """Upload files to s3 with splitting into batches of the specified size. Returns list of keys for each batch."""
    batch_size = config.BATCH_SIZE
    log.info(f"Configured batch size: {batch_size}")
    total_batches = math.ceil(len(account_to_json) / batch_size)
    log.info(f"Total batches to process: {total_batches}")
    batch_keys = []

    def process_batch(batch_number: int):
        start_index = batch_number * batch_size
        end_index = start_index + batch_size
        batch = account_to_json[start_index:end_index]
        log.info(f"Processing batch {batch_number + 1} of {total_batches}...")
        log.info(f"Batch size: {len(batch)}")

        min_contract_id = _get_min_contract_id(batch)
        max_contract_id = _get_max_contract_id(batch)
        log.info(f"Batch min_contract_id: {min_contract_id}, max_contract_id: {max_contract_id}")

        batch_key = f"{key_prefix}/batch_{batch_number + 1}_{min_contract_id}_{max_contract_id}"
        log.info(f"Generated batch key: {batch_key}")

        _upload_batch_to_s3(batch, batch_key, bucket)
        log.info(f"Batch {batch_number + 1} uploaded successfully.")

        return batch_key

    with ThreadPoolExecutor(max_workers=6) as executor:
        future_to_batch = {executor.submit(process_batch, batch_number): batch_number for batch_number in
                           range(total_batches)}
        for future in as_completed(future_to_batch):
            batch_keys.append(future.result())

    log.info("All batches processed successfully.")
    return batch_keys