import concurrent
import json
import logging
import time
from concurrent.futures import ThreadPoolExecutor

import config
from oa_contract_ingest_to_abacus.connectors import aws_lambda

log = logging.getLogger(__name__)


def run_lambda_for_batches(
        batch_keys: list[str],
        import_user: str,
        jira_ticket: str,
        lambda_function_name: str,
) -> list[dict]:
    tasks = []
    for batch_key in batch_keys:
        tasks.append((
            lambda_function_name,
            {
                'directory': batch_key,
                'overwrite': 1,
                'import_user': import_user,
                'jira_ticket': jira_ticket,
                'skip_account_import': 1,
                # 'skip_contract_import': 1,
            }
        ))
    log.info(f"Created {len(tasks)} tasks for batch processing.")
    return _run_parallel_lambdas(tasks, config.LAMBDA_MAX_PARALLEL)


def _invoke_lambda(function_name: str, payload: dict, max_retries: int = 10, retry_delay: int = 60) -> dict:
    """Invoke with retry."""
    attempt = 0

    while attempt <= max_retries:
        try:
            start_time = time.time()
            log.info(f"Invoking Lambda function '{function_name}' with payload: {payload}")
            response = aws_lambda.invoke_lambda(function_name, json.dumps(payload))
            end_time = time.time()
            execution_time = end_time - start_time
            log.info(f"Received response from Lambda function '{function_name}' and payload {payload}: {response}")
            log.info(f"Execution time for Lambda function without retries '{function_name}': {execution_time:.2f} seconds")

            return json.loads(response)
        except Exception as e:
            attempt += 1
            if attempt > max_retries:
                log.error(f"Maximum retry attempts ({max_retries}) reached for Lambda function '{function_name}'.")
                raise e
            else:
                log.warning(f"Exception encountered: {str(e)}. Retrying attempt {attempt} "
                            f"out of {max_retries} in {retry_delay} seconds...")
                time.sleep(retry_delay)

    raise RuntimeError(f"Function '{function_name}' failed to invoke after {max_retries} retries with payload {payload}.")


def _run_parallel_lambdas(tasks: list[tuple[str, dict]], max_workers) -> list[dict | None]:
    log.info(f"Starting parallel execution of {len(tasks)} tasks with max_workers={max_workers}.")
    results: list[dict | None] = [None] * len(tasks)

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        # Submit all tasks
        future_to_index = {
            executor.submit(_invoke_lambda, fn, payload): i
            for i, (fn, payload) in enumerate(tasks)
        }

        for future in concurrent.futures.as_completed(future_to_index):
            index = future_to_index[future]
            try:
                results[index] = future.result()
                log.info(f"Task {index} completed successfully.")
            except Exception as e:
                results[index] = {"error": str(e)}
                log.info(f"Task {index} failed with error: {e}")

    log.info("All tasks have been processed.")
    return results