import hashlib
import json
import traceback
from tqdm import asyncio as tqdmasyncio
import os
import signal
import sys
import logging
import time  # Added for timestamping logs
import asyncio  # Add asyncio import

from augmentoolkit.generation_functions.pipeline_step_class import PipelineStep
from filelock import FileLock

# Try importing pyfiglet for banners, provide fallback
try:
    import pyfiglet

    def _print_banner(text):  # Use underscore to indicate internal helper
        print(pyfiglet.figlet_format(text, font="standard"))

except ImportError:
    print(
        "Warning: pyfiglet not installed. Install with 'pip install pyfiglet' for formatted banners."
    )

    def _print_banner(text):
        print(f"\n{'='*10} {text} {'='*10}\n")


# Helper function similar to the one in PipelineStep
def filter_out_nonpresent_keys(item_dict, key_to_check="judgement"):
    # print(f"[DEBUG] filter_out_nonpresent_keys: Starting filter for key '{key_to_check}' on dict with {len(item_dict)} items.") # DEBUG
    keys_to_remove = []
    for key, value in item_dict.items():
        # Ensure value is a dictionary before checking for the key
        if isinstance(value, dict) and key_to_check not in value:
            # print(f"[DEBUG] filter_out_nonpresent_keys: Marking key '{key}' for removal (missing '{key_to_check}').") # DEBUG
            keys_to_remove.append(key)
        elif not isinstance(
            value, dict
        ):  # Handle cases where value might not be a dict unexpectedly
            print(
                f"Warning: Expected dictionary for key '{key}' in filter_out_nonpresent_keys, but got {type(value)}. Skipping."
            )
            # Optionally remove if malformed data should be cleaned
            # keys_to_remove.append(key)

    if keys_to_remove:
        print(
            f"Filtering out {len(keys_to_remove)} items missing the key '{key_to_check}'"
        )
        for key in keys_to_remove:
            del item_dict[key]
    # print(f"[DEBUG] filter_out_nonpresent_keys: Finished filtering. {len(item_dict)} items remain.") # DEBUG


class OneToManyStep(
    PipelineStep
):  # A way of switching from one list of items to a new list of many items (each part of a group generated from the input items). Like going from a small list of source paragraphs to a larger list of question/answer pairs to iterate over.
    # fundamentally, it is easier if we carry over all previous information. Like, all previous keys. Yeah. Let's start with that they can subclass it if they want.
    def __init__(
        self,
        prompt_path,
        sampling_params,
        output_file,
        output_processor,
        result_key,
        max_retries,
        details_key,
        log_full_outputs,
        input_file,
        validation_function=lambda x, y: {"result": True, "message": "default message"},
        regex=r"(.*)",
        input_processor=lambda x: (x, {}),
        **kwargs,
    ):
        # print(f"[DEBUG] OneToManyStep.__init__: Initializing step for output file '{output_file}'.") # DEBUG
        super().__init__(
            prompt_path=prompt_path,
            regex=regex,
            sampling_params=sampling_params,
            output_file=output_file,
            output_processor=output_processor,
            result_key=result_key,
            max_retries=max_retries,
            log_full_outputs=log_full_outputs,
            validation_function=validation_function,
            details_key=details_key,
            input_processor=input_processor,
            **kwargs,
        )
        self.input_file = input_file
        # print(f"[DEBUG] OneToManyStep.__init__: Initialization complete for '{output_file}'.") # DEBUG

    async def read_previous_output(self, key, output_dict):
        # print(f"[DEBUG {time.time():.2f}] read_previous_output: Checking for key '{key}' or prefix '{key}-' in output_dict (size {len(output_dict)}).") # DEBUG
        # Check if any keys start with this key followed by separator in the in-memory dict
        key_prefix = str(key) + "-"
        matching_entries = [k for k in output_dict.keys() if k.startswith(key_prefix)]
        if matching_entries:
            # print(f"[DEBUG {time.time():.2f}] read_previous_output: Found {len(matching_entries)} existing entries for key prefix '{key}-'. Skipping generation.") # DEBUG
            # No need to load from file, data should already be in output_dict if loaded
            # print(f"Found existing entries for key {key} in memory: {matching_entries}")
            return True  # Indicate that output for this key exists
        # print(f"[DEBUG {time.time():.2f}] read_previous_output: No existing entries found for key prefix '{key}-'. Proceeding with generation.") # DEBUG
        return False

    def save(
        self,
        result=None,
        key=None,
        output_dict=None,
        input_data=None,
        full_output=None,
        full_response=None,
        full_input=None,
        include_details=False,
        completion_mode=False,
    ):
        # print(f"[DEBUG {time.time():.2f}] save: Starting save for key '{key}'. Result type: {type(result)}, Result length: {len(result) if isinstance(result, list) else 'N/A'}.") # DEBUG
        # Removed output_dir, input_dict args as they are handled by execute_pipeline
        # print("\n=== Starting save() function (In-Memory Update) ===")
        # print(f"Key: {key}")
        # print(f"Result type: {type(result)}")
        # print(f"Result length: {len(result)}")
        # print(f"Include details: {include_details}")

        assert isinstance(
            result, list
        ), f"Result must be a list, got {type(result)}"  # needed for one-to-many
        # print(f"[DEBUG {time.time():.2f}] save: Assertion passed, result is a list.") # DEBUG

        # Removed output_path and related file operations

        if self.log_full_outputs:
            # print(f"[DEBUG {time.time():.2f}] save: log_full_outputs is True. Logging full outputs (logic omitted for brevity).") # DEBUG
            # print("\nLogging full outputs:") # Keep logging if needed
            if full_output:
                # ... (logging logic remains the same) ...
                pass  # Pass for brevity, assume logging logic is unchanged

        # print("\n=== Processing result items (In-Memory Update) ===")
        details_payload = None
        if include_details:
            # print(f"[DEBUG {time.time():.2f}] save: include_details is True. Preparing details payload.") # DEBUG
            # Prepare details payload once, assign it to each subitem
            details_payload = [
                {
                    "full_response": full_response,
                    "full_input": full_input,
                    "completion_mode": completion_mode,
                    "result": result,
                }
            ]
            # print("Details payload prepared:")
            # print(details_payload)

        for idx, subitem in enumerate(result):
            # print(f"[DEBUG {time.time():.2f}] save: Processing item {idx} for original key '{key}'. Subitem: {str(subitem)[:50]}...") # DEBUG
            # print(f"\nProcessing item {idx}:")
            # print(f"Subitem: {str(subitem)[:100]}...")

            # Hash the idx and subitem using hashlib
            hash_input = f"{str(subitem)}"
            item_hash = hashlib.sha256(hash_input.encode()).hexdigest()
            item_key = f"{key}-{idx}-{item_hash}"
            # print(f"[DEBUG {time.time():.2f}] save: Generated item_key: {item_key}") # DEBUG
            # print(f"Generated item_key: {item_key}")

            # Create new entry with ALL original data plus the new result
            output_dict[str(item_key)] = (
                input_data.copy()
            )  # Start with a copy of the original input
            # print(f"[DEBUG {time.time():.2f}] save: Copied input_data to new key '{item_key}'.") # DEBUG
            output_dict[str(item_key)][
                self.result_key
            ] = subitem  # Add the specific subitem result
            # print(f"[DEBUG {time.time():.2f}] save: Added subitem result under key '{self.result_key}' for item_key '{item_key}'.") # DEBUG
            # print(f"Added result to output_dict with key: {self.result_key}")

            if include_details and details_payload:
                output_dict[str(item_key)][
                    self.details_key
                ] = details_payload  # Add details payload
                # print(f"[DEBUG {time.time():.2f}] save: Added details payload under key '{self.details_key}' for item_key '{item_key}'.") # DEBUG
                # print(f"Added details to output_dict with key: {self.details_key}")

            # Removed file writing logic for each item

        # Removed logic that modified the input_file/input_dict for details

        # print(f"[DEBUG {time.time():.2f}] save: Finished processing {len(result)} items for key '{key}'. output_dict size now {len(output_dict)}.") # DEBUG
        # print("\n=== save() function completed (In-Memory Update) ===\n")
        # Save now returns None as it modifies output_dict in place
        return (
            None  # Or return result if needed elsewhere, but primary effect is mutation
        )

    def load_dataset(self, output_dict, output_dir):
        output_path = self.make_output_path(output_dir)
        # print(f"[DEBUG {time.time():.2f}] load_dataset: Attempting to load dataset for '{self.output_file}' from '{output_path}'.") # DEBUG
        if os.path.exists(output_path):
            # Use file lock for reading, although less critical than writing
            with open(output_path, "r", encoding="utf-8") as f:
                try:
                    print(
                        f"Loading existing output file for {self.output_file} at {output_path}"
                    )
                    file_contents = json.load(f)
                except json.JSONDecodeError:
                    print("\n" + "!" * 80)
                    print(f"WARNING: JSON DECODE ERROR in file {output_path}")
                    print(
                        "The file exists but contains invalid JSON. Starting with an empty dictionary instead."
                    )
                    print(
                        "This may indicate data corruption or an incomplete write operation."
                    )
                    print("!" * 80 + "\n")
                    file_contents = {}  # Start fresh if decode fails

            assert isinstance(
                file_contents, dict
            ), f"Expected file_contents to be a dictionary, but got {type(file_contents).__name__}."
            # print(f"[DEBUG {time.time():.2f}] load_dataset: File read successfully. Loaded {len(file_contents)} items. Updating output_dict.") # DEBUG
            output_dict.update(file_contents)  # Load into the provided dict
            print(f"Loaded {len(file_contents)} items from {output_path}")
        else:
            # print(f"[DEBUG {time.time():.2f}] load_dataset: Output file '{output_path}' not found. Starting with an empty dataset.") # DEBUG
            print(
                f"Output file {output_path} not found. Starting with an empty dataset."
            )
            # Ensure directory exists even if file doesn't
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            # print(f"[DEBUG {time.time():.2f}] load_dataset: Ensured directory '{os.path.dirname(output_path)}' exists.") # DEBUG

    def save_dataset(self, output_dict, output_dir):
        output_path = self.make_output_path(output_dir=output_dir)
        # print(f"\n[DEBUG {time.time():.2f}] save_dataset: Attempting to save dataset for '{self.output_file}' ({len(output_dict)} items) to '{output_path}'.") # DEBUG
        print(f"\nAttempting to save dataset... {output_path}")

        # Robust save with interrupt handling (adapted from PipelineStep)
        interrupt_count = 0
        max_interrupts = 5
        original_sigint_handler = signal.getsignal(signal.SIGINT)

        def _robust_save_interrupt_handler(sig, frame):
            nonlocal interrupt_count
            interrupt_count += 1
            if interrupt_count == 1:
                _print_banner("SAVE WARNING")
                # print(f"[DEBUG {time.time():.2f}] _robust_save_interrupt_handler: First interrupt during save for '{self.output_file}'.") # DEBUG
                print(
                    f"\nWARNING: Interrupt detected during critical dataset save for step '{self.output_file}'!",
                    file=sys.stderr,
                )
                print(
                    f"Press Ctrl+C {max_interrupts - interrupt_count} more times consecutively to force exit (DATA LOSS RISK!).",
                    file=sys.stderr,
                )
            elif interrupt_count < max_interrupts:
                print(
                    f"\nInterrupt {interrupt_count}/{max_interrupts}. Press Ctrl+C {max_interrupts - interrupt_count} more times to force exit.",
                    file=sys.stderr,
                )
                # print(f"[DEBUG {time.time():.2f}] _robust_save_interrupt_handler: Interrupt {interrupt_count}/{max_interrupts} during save for '{self.output_file}'.") # DEBUG
            else:
                print(
                    f"\n{max_interrupts} consecutive interrupts detected during save. Forcing exit.",
                    file=sys.stderr,
                )
                # print(f"[DEBUG {time.time():.2f}] _robust_save_interrupt_handler: Max interrupts reached for '{self.output_file}'. Forcing exit.") # DEBUG
                signal.signal(
                    signal.SIGINT, original_sigint_handler
                )  # Restore before exit
                sys.exit(130)  # Exit code 130 for SIGINT

        save_successful = False
        try:
            # print(f"[DEBUG {time.time():.2f}] save_dataset: Setting SIGINT handler for robust save.") # DEBUG
            signal.signal(
                signal.SIGINT, _robust_save_interrupt_handler
            )  # Set custom handler

            # Ensure directory exists
            # print(f"[DEBUG {time.time():.2f}] save_dataset: Ensuring directory '{os.path.dirname(output_path)}' exists.") # DEBUG
            os.makedirs(os.path.dirname(output_path), exist_ok=True)

            # Use file lock for writing
            # print(f"[DEBUG {time.time():.2f}] save_dataset: Lock acquired. Starting atomic write.") # DEBUG
            # Atomic write using temp file
            temp_output_path = (
                f"{output_path}.temp_{os.getpid()}"  # Add PID for extra safety
            )
            # print(f"[DEBUG {time.time():.2f}] save_dataset: Writing to temp file: '{temp_output_path}'.") # DEBUG
            try:
                with open(temp_output_path, "w", encoding="utf-8") as f:
                    json.dump(
                        output_dict, f, indent=None, ensure_ascii=True
                    )  # Use ensure_ascii=True for cross-platform compatibility
                # print(f"[DEBUG {time.time():.2f}] save_dataset: Temp file write complete. Replacing original file.") # DEBUG

                os.replace(temp_output_path, output_path)  # Atomic rename/replace
                # print(f"[DEBUG {time.time():.2f}] save_dataset: Atomic replace successful.") # DEBUG
                save_successful = True
            except Exception as write_error:
                # print(f"[DEBUG {time.time():.2f}] save_dataset: ERROR during file write/replace: {write_error}") # DEBUG
                print(
                    f"\nERROR during file write/replace operation: {write_error}",
                    file=sys.stderr,
                )
                traceback.print_exc()
                # Clean up temp file if replace failed
                if os.path.exists(temp_output_path):
                    # print(f"[DEBUG {time.time():.2f}] save_dataset: Attempting to remove failed temp file '{temp_output_path}'.") # DEBUG
                    try:
                        os.remove(temp_output_path)
                    except OSError as remove_error:
                        print(
                            f"Error removing temp file {temp_output_path}: {remove_error}",
                            file=sys.stderr,
                        )
                raise  # Re-raise the error after attempting cleanup
            finally:
                # Ensure temp file is removed if it somehow still exists after a successful replace (unlikely but safe)
                if save_successful and os.path.exists(temp_output_path):
                    # print(f"[DEBUG {time.time():.2f}] save_dataset: Ensuring temp file '{temp_output_path}' is removed after successful save.") # DEBUG
                    try:
                        os.remove(temp_output_path)
                    except OSError:
                        pass  # Ignore error if removal fails post-success
            # print(f"[DEBUG {time.time():.2f}] save_dataset: Lock released.") # DEBUG

        except Exception as e:
            # print(f"[DEBUG {time.time():.2f}] save_dataset: Exception during save process: {e}") # DEBUG
            print(f"\nAn error occurred during the save process: {e}", file=sys.stderr)
            # The finally block handles handler restoration.
            # Consider re-raising if the pipeline should halt on save failure.
            # raise e
        finally:
            # print(f"[DEBUG {time.time():.2f}] save_dataset: Restoring original SIGINT handler.") # DEBUG
            signal.signal(
                signal.SIGINT, original_sigint_handler
            )  # ALWAYS restore handler
            if save_successful:
                # print(f"[DEBUG {time.time():.2f}] save_dataset: Save successful for '{output_path}'.") # DEBUG
                print(
                    f"Dataset for '{self.output_file}' saved successfully to {output_path}."
                )
            elif interrupt_count < max_interrupts:
                # print(f"[DEBUG {time.time():.2f}] save_dataset: Save may be incomplete for '{output_path}' due to error or interrupt.") # DEBUG
                print(
                    f"Dataset save for '{self.output_file}' may be incomplete due to an error or interruption.",
                    file=sys.stderr,
                )
            # Reset counter is implicit as it's local to the next call
            # print(f"[DEBUG {time.time():.2f}] save_dataset: Save process finished.") # DEBUG

    async def run(
        self,
        key=None,
        input_data=None,
        engine_wrapper=None,
        default_prompt_folder=None,
        prompt_folder=None,
        completion_mode=False,
        use_stop=True,
        output_dict=None,
        include_details=False,
        **kwargs,
    ):
        # print(f"[DEBUG {time.time():.2f}] run: Starting for key '{key}'. Completion mode: {completion_mode}.") # DEBUG
        full_prompt_path = (
            self.prompt_path + ".yaml"
            if not completion_mode
            else self.prompt_path + ".txt"
        )
        # print(f"[DEBUG {time.time():.2f}] run: Prompt path: '{full_prompt_path}'.") # DEBUG

        # Check existing entries IN MEMORY
        # print(f"[DEBUG {time.time():.2f}] run: Checking previous output for key '{key}'.") # DEBUG
        if await self.read_previous_output(key, output_dict):  # Pass output_dict
            # print(f"[DEBUG {time.time():.2f}] run: Previous output found for key '{key}'. Skipping.") # DEBUG
            return  # Skip generation if already processed
        # print(f"[DEBUG {time.time():.2f}] run: No previous output found for key '{key}'. Proceeding.") # DEBUG

        # print(f"[DEBUG {time.time():.2f}] run: Processing input data for key '{key}'.") # DEBUG
        if asyncio.iscoroutinefunction(self.process_input_data):
            processed_data, additional_kwargs = await self.process_input_data(input_data)
        else:
            processed_data, additional_kwargs = self.process_input_data(input_data)
        # print(f"[DEBUG {time.time():.2f}] run: Input data processed for key '{key}'.") # DEBUG

        error_message = ""
        complete = False
        retries_left = self.max_retries
        # print(f"[DEBUG {time.time():.2f}] run: Starting generation loop for key '{key}'. Max retries: {retries_left}.") # DEBUG
        while not complete and retries_left > 0:
            # print(f"[DEBUG {time.time():.2f}] run: Attempting generation for key '{key}'. Retries left: {retries_left}.") # DEBUG
            try:
                # print(f"[DEBUG {time.time():.2f}] run: Calling generate_data for key '{key}'.") # DEBUG
                start_time = time.time()  # DEBUG TIME
                result, full_output, full_response, full_input = (
                    await self.generate_data(
                        processed_data=processed_data,
                        engine_wrapper=engine_wrapper,
                        prompt_path=full_prompt_path,
                        prompt_folder=prompt_folder,
                        default_prompt_folder=default_prompt_folder,
                        completion_mode=completion_mode,
                        use_stop=use_stop,
                        error_message=error_message,
                        **kwargs,
                        **additional_kwargs,
                    )
                )
                end_time = time.time()  # DEBUG TIME
                # print(f"[DEBUG {end_time:.2f}] run: generate_data call for key '{key}' completed in {end_time - start_time:.2f} seconds.") # DEBUG
                # Ensure result is a list for OneToMany validation
                if not isinstance(result, list):
                    # print(f"[DEBUG {time.time():.2f}] run: WARNING - Expected list result from generate_data for key {key}, got {type(result)}. Raising TypeError.") # DEBUG
                    print(
                        f"Warning: Expected list result from generate_data for key {key}, got {type(result)}. Treating as failed attempt."
                    )
                    raise TypeError("Generator did not return a list")  # Trigger retry

                # print(f"[DEBUG {time.time():.2f}] run: Generation successful for key '{key}'. Result is a list of length {len(result)}. Validating...") # DEBUG
                validation_result = self.validation_function(result, input_data)
                if validation_result["result"]:
                    # print(f"[DEBUG {time.time():.2f}] run: Validation successful for key '{key}'.") # DEBUG
                    complete = True
                else:
                    # print(f"[DEBUG {time.time():.2f}] run: Validation FAILED for key {key}.") # DEBUG
                    error_message = validation_result["message"]
                    print(f"Validation failed for key {key}")

            except Exception as e:
                # print(f"[DEBUG {time.time():.2f}] run: Exception during generation/validation for key '{key}': {type(e).__name__}: {e}") # DEBUG
                print(e)
                error_message = str(e)
                traceback.print_exc()
            retries_left -= 1
            if not complete and retries_left > 0:
                # print(f"[DEBUG {time.time():.2f}] run: Retrying generation for key '{key}'. {retries_left} retries remaining.") # DEBUG
                await asyncio.sleep(1)  # Optional: Add a small delay before retrying
            elif not complete:
                # print(f"[DEBUG {time.time():.2f}] run: Generation failed for key '{key}' after exhausting retries.") # DEBUG
                pass  # Added pass to fix potential indentation error
        if not complete:  # consider raising here and catching in the actual pipeline.
            # print(f"[DEBUG {time.time():.2f}] run: Generation ultimately failed for key '{key}'. Returning None.") # DEBUG
            return

        # print(f"[DEBUG {time.time():.2f}] run: Generation complete for key '{key}'. Calling save.") # DEBUG
        save_result = self.save(
            result=result,
            key=key,
            output_dict=output_dict,  # Pass the shared dictionary
            input_data=input_data,
            full_output=full_output,
            full_response=full_response,
            full_input=full_input,
            include_details=include_details,
            completion_mode=completion_mode,
        )
        # print(f"[DEBUG {time.time():.2f}] run: Save call completed for key '{key}'. Returning result: {type(save_result)}") # DEBUG
        return save_result

    async def execute_pipeline(
        self,
        input_dict={},
        engine_wrapper=None,
        rtwl=None,
        default_prompt_folder=None,
        prompt_folder=None,
        output_dir=None,
        completion_mode=None,
        use_stop=None,
        include_details=False,
        **kwargs,
    ):
        # print(f"[DEBUG {time.time():.2f}] execute_pipeline: Starting execution for step '{self.output_file}'. Input dict size: {len(input_dict)}.") # DEBUG
        output_dict = {}  # Initialize the dictionary for this step's results
        # print(f"[DEBUG {time.time():.2f}] execute_pipeline: Loading dataset for '{self.output_file}'.") # DEBUG
        self.load_dataset(output_dict, output_dir)  # Load existing data at the start
        # print(f"[DEBUG {time.time():.2f}] execute_pipeline: Dataset loaded. Current output_dict size: {len(output_dict)}.") # DEBUG

        try:
            # Note: input_dict here refers to the *previous* step's output
            # print(f"[DEBUG {time.time():.2f}] execute_pipeline: Preparing generation tasks from input_dict (size {len(input_dict)}).") # DEBUG
            data_generations_tasks = []
            for i, (key, value) in enumerate(input_dict.items()):
                # print(f"[DEBUG {time.time():.2f}] execute_pipeline: Creating task {i+1}/{len(input_dict)} for key '{key}'.") # DEBUG
                task = self.run(
                    input_data=value,
                    engine_wrapper=engine_wrapper,
                    key=key,
                    default_prompt_folder=default_prompt_folder,
                    prompt_folder=prompt_folder,
                    completion_mode=completion_mode,
                    use_stop=use_stop,
                    output_dict=output_dict,  # Pass the shared dict to run/save
                    include_details=include_details,
                    **kwargs,
                )
                data_generations_tasks.append(task)
            # data_generations_tasks = [
            #     self.run(
            #         input_data=value,
            #         engine_wrapper=engine_wrapper,
            #         key=key,
            #         default_prompt_folder=default_prompt_folder,
            #         prompt_folder=prompt_folder,
            #         completion_mode=completion_mode,
            #         use_stop=use_stop,
            #         output_dict=output_dict, # Pass the shared dict to run/save
            #         include_details=include_details,
            #         **kwargs
            #     ) for key, value in input_dict.items() # Iterate over inputs from previous step
            # ]

            if not data_generations_tasks:
                # print(f"[DEBUG {time.time():.2f}] execute_pipeline: Warning - No tasks created for step '{self.output_file}'.") # DEBUG
                print("Warning: No tasks to run for this step.")

            # print(f"[DEBUG {time.time():.2f}] execute_pipeline: Created {len(data_generations_tasks)} generation tasks. Wrapping with Rate-Limited Wrapper (rtwl).") # DEBUG
            coroutines = [rtwl(task) for task in data_generations_tasks]
            # print(f"[DEBUG {time.time():.2f}] execute_pipeline: Wrapped {len(coroutines)} tasks. Starting processing with tqdm.") # DEBUG
            print(f"Processing {len(coroutines)} items...")
            TASK_TIMEOUT_SECONDS = 600  # 10 minutes timeout
            processed_count = 0
            total_tasks = len(coroutines)
            for future in tqdmasyncio.tqdm(
                tqdmasyncio.tqdm.as_completed(coroutines),
                total=len(coroutines),
                desc=f"Step: {self.output_file}",
            ):
                processed_count += 1  # DEBUG
                # print(f"[DEBUG {time.time():.2f}] execute_pipeline: Awaiting future {processed_count}/{len(coroutines)}.") # DEBUG - Can be very verbose
                try:
                    # Wrap the await future with a timeout
                    await asyncio.wait_for(future, timeout=TASK_TIMEOUT_SECONDS)
                    # print(f"[DEBUG {time.time():.2f}] execute_pipeline: Future {processed_count}/{len(coroutines)} completed successfully.") # DEBUG
                except asyncio.TimeoutError:
                    print(
                        f"\nWARNING: Task {processed_count}/{total_tasks} for step '{self.output_file}' timed out after {TASK_TIMEOUT_SECONDS} seconds.",
                        file=sys.stderr,
                    )
                    # Task is cancelled by wait_for, loop continues
                except Exception as task_exc:
                    # print(f"[DEBUG {time.time():.2f}] execute_pipeline: Error in awaited task {processed_count}/{len(coroutines)}: {task_exc}") # DEBUG
                    # Log exceptions from individual run tasks if needed, but allow others to continue
                    print(f"\nError in awaited task: {task_exc}")
                    traceback.print_exc()  # Optionally print traceback for task errors
            # print(f"[DEBUG {time.time():.2f}] execute_pipeline: Finished processing all {len(coroutines)} futures.") # DEBUG

            # Filter out entries where the generation might have failed and didn't add the result key
            # This is crucial for OneToMany as failed 'run' won't add keys to output_dict
            # but we still need to ensure consistency if partial results were loaded.
            # However, the primary filtering happens because failed `run` calls simply don't add anything to `output_dict`.
            # A check after loading might be useful for corrupted files, but less so for runtime failures.
            # Consider if filtering is needed based on how `load_dataset` handles errors.
            # print(f"[DEBUG {time.time():.2f}] execute_pipeline: Optional filtering step (commented out).") # DEBUG
            # filter_out_nonpresent_keys(output_dict, key_to_check=self.result_key) # Optional: clean loaded data

        except Exception as e:
            # print(f"[DEBUG {time.time():.2f}] execute_pipeline: CRITICAL EXCEPTION during execution for step '{self.output_file}': {e}") # DEBUG
            print(
                f"\nCritical exception during pipeline execution for step {self.output_file}: {e}"
            )
            traceback.print_exc()
            # The finally block will still attempt to save. Re-raise if pipeline should stop.
            raise e
        finally:
            final_count = len(output_dict)
            # print(f"[DEBUG {time.time():.2f}] execute_pipeline: FINALLY block reached. Final item count: {final_count}. Saving dataset.") # DEBUG
            print(
                f"\nSaving results for step '{self.output_file}'. Final item count: {final_count}"
            )
            # Robust save called in finally block
            filter_out_nonpresent_keys(output_dict, key_to_check=self.result_key)
            self.save_dataset(output_dict, output_dir)
            # print(f"[DEBUG {time.time():.2f}] execute_pipeline: FINALLY block complete after save_dataset call.") # DEBUG

        # print(f"[DEBUG {time.time():.2f}] execute_pipeline: Execution finished for step '{self.output_file}'. Returning output_dict with {len(output_dict)} items.") # DEBUG
        return output_dict  # Return the completed dictionary