#!/usr/bin/env python3

from __future__ import annotations

import time
import multiprocessing
import pytest
import threading
import ctypes
import queue
from dataclasses import dataclass
from queue import Empty

from sgn.sinks import NullSink
from sgn.sources import SignalEOS
from sgn.subprocess import (
    Parallelize,
    _ParallelizeBase,
    ParallelizeTransformElement,
    ParallelizeSinkElement,
    ParallelizeSourceElement,
    WorkerContext,
    _worker_wrapper_function,
)
from sgn.base import SourceElement, Frame
from sgn.apps import Pipeline


# Add fixture for test isolation and cleanup
@pytest.fixture(scope="function", autouse=True)
def clean_subprocess_state(monkeypatch):
    """Reset SubProcess state between tests."""
    import gc

    # Create clean copies for this test
    instance_list_copy = []
    shm_list_copy = []

    # Keep track of the original threading default to restore it
    original_threading_default = Parallelize.use_threading_default

    # Clean up any existing shared memory
    try:
        shm = multiprocessing.shared_memory.SharedMemory(name="shared_data")
        shm.unlink()
        shm.close()
    except (FileNotFoundError, ValueError):
        pass

    # Monkeypatch the class attributes to use our clean copies
    monkeypatch.setattr(Parallelize, "instance_list", instance_list_copy)
    monkeypatch.setattr(Parallelize, "shm_list", shm_list_copy)

    # Reset use_threading_default to False (multiprocessing mode)
    monkeypatch.setattr(Parallelize, "use_threading_default", False)

    # Ensure clean state
    gc.collect()

    # Let the test run
    yield

    # Teardown - clean up resources
    for p in instance_list_copy:
        if hasattr(p, "worker") and p.worker and hasattr(p.worker, "is_alive"):
            try:
                if p.worker.is_alive():
                    if hasattr(p.worker, "terminate"):
                        p.worker.terminate()
                    p.worker.join(timeout=0.1)
            except (ValueError, RuntimeError):
                pass

    # Clean up all shared memory
    for d in shm_list_copy:
        try:
            if "name" in d:
                multiprocessing.shared_memory.SharedMemory(name=d["name"]).unlink()
        except (FileNotFoundError, ValueError):
            pass

    # Explicitly clean shared_data
    try:
        shm = multiprocessing.shared_memory.SharedMemory(name="shared_data")
        shm.unlink()
        shm.close()
    except (FileNotFoundError, ValueError):
        pass

    # Explicitly clean test_duplicate
    try:
        shm = multiprocessing.shared_memory.SharedMemory(name="test_duplicate")
        shm.unlink()
        shm.close()
    except (FileNotFoundError, ValueError):
        pass

    # Restore the original threading default
    monkeypatch.setattr(
        Parallelize, "use_threading_default", original_threading_default
    )

    gc.collect()


def get_address(buffer):
    address = ctypes.addressof(ctypes.c_char.from_buffer(buffer))
    return address


#
# Simple test elements for basic subprocess functionality
#
class MySourceClass(SourceElement, SignalEOS):
    """A simple source class that just sends an EOS frame."""

    def new(self, pad):
        return Frame(data=None, EOS=True)


@dataclass
class MySinkClass(ParallelizeSinkElement):
    """A sink class that does minimal processing for testing."""

    def __post_init__(self):
        super().__post_init__()

    def pull(self, pad, frame):
        if frame.EOS:
            self.mark_eos(pad)
        if self.at_eos and not self.terminated.is_set():
            self.in_queue.put(frame)
            self.sub_process_shutdown(30)  # Increased timeout for CI environments

    def worker_process(self, context):
        context.output_queue.put(None)
        try:
            context.input_queue.get(timeout=0.1)
        except queue.Empty:
            pass


@dataclass
class MyTransformClass(ParallelizeTransformElement):
    """A transform class that runs in a separate process."""

    def __post_init__(self):
        super().__post_init__()
        assert len(self.sink_pad_names) == 1 and len(self.source_pad_names) == 1
        self.at_eos = False
        self.frame_list = []

    def pull(self, pad, frame):
        self.in_queue.put(frame)
        if frame.EOS and not self.terminated.is_set():
            self.at_eos = True
            self.frame_list = self.sub_process_shutdown(10)

    def worker_process(self, context):
        # access some shared memory - there is only one
        # Just access it to verify it exists, but don't need to use it
        if context.shared_memory:
            _ = context.shared_memory[0]["shm"]
        try:
            frame = context.input_queue.get(timeout=0.1)
            if frame:
                context.output_queue.put(frame)
        except queue.Empty:
            pass

    def new(self, pad):
        return self.frame_list[0]


#
# Elements for testing concurrency modes
#
class NumberSource(SourceElement, SignalEOS):
    """A simple source that generates a sequence of numbers."""

    def __init__(self, count=5, **kwargs):
        super().__init__(**kwargs)
        self.count = count
        self.current = 0

    def new(self, pad):
        if self.current >= self.count:
            return Frame(data=None, EOS=True)
        self.current += 1
        return Frame(data=self.current, EOS=False)


# Source element test classes
@dataclass
class SimpleThreadedSource(ParallelizeSourceElement):
    """A simple source element that generates sequential numbers using a thread."""

    _use_threading_override: bool = True
    count: int = 3

    def __post_init__(self):
        super().__post_init__()
        # Track EOS status per output pad
        self.pad_eos_sent = {pad.name: False for pad in self.source_pads}
        # count will be automatically passed to worker_process
        # Pre-populate result queue for testing
        self.results = []

    def new(self, pad):
        """Get the next frame for the given pad."""
        # If we've already marked this pad as EOS, keep returning EOS frames
        if self.pad_eos_sent.get(pad.name, False):
            return Frame(data=None, EOS=True)

        try:
            # Try to get data from the queue with a very short timeout
            data = self.out_queue.get(timeout=0.05)

            # None signals EOS
            if data is None:
                self.pad_eos_sent[pad.name] = True
                # If all pads have reached EOS, set the global EOS flag
                if all(self.pad_eos_sent.values()):
                    self.at_eos = True
                return Frame(data=None, EOS=True)

            # Store result for test verification
            self.results.append(data)
            # Return regular data frame
            return Frame(data=data)

        except Empty:
            # If queue is empty, return empty frame
            return Frame(data=None)

    def worker_process(self, context: WorkerContext, count: int) -> None:
        """Generate sequential numbers and send to the main thread."""
        # Send count number of items
        for i in range(1, count + 1):
            # Check if we should stop
            if context.should_stop():
                break

            # Send the number
            context.output_queue.put(i)
            # Minimal delay to avoid flooding the queue while keeping tests fast
            time.sleep(0.001)

        # Signal end of stream
        context.output_queue.put(None)


@dataclass
class SimpleProcessSource(ParallelizeSourceElement):
    """A simple source element that generates squared numbers using a process."""

    _use_threading_override: bool = False  # Use multiprocessing
    count: int = 3

    def __post_init__(self):
        super().__post_init__()
        # Track EOS status per output pad
        self.pad_eos_sent = {pad.name: False for pad in self.source_pads}
        # count will be automatically passed to worker_process
        # Pre-populate result queue for testing
        self.results = []

    def new(self, pad):
        """Get the next frame for the given pad."""
        # If we've already marked this pad as EOS, keep returning EOS frames
        if self.pad_eos_sent.get(pad.name, False):
            return Frame(data=None, EOS=True)

        try:
            # Try to get data from the queue with a very short timeout
            data = self.out_queue.get(timeout=0.05)

            # None signals EOS
            if data is None:
                self.pad_eos_sent[pad.name] = True
                # If all pads have reached EOS, set the global EOS flag
                if all(self.pad_eos_sent.values()):
                    self.at_eos = True
                return Frame(data=None, EOS=True)

            # Store result for test verification
            self.results.append(data)
            # Return regular data frame
            return Frame(data=data)

        except Empty:
            # If queue is empty, return empty frame
            return Frame(data=None)

    def worker_process(self, context: WorkerContext, count: int) -> None:
        """Generate squared numbers and send to the main process."""
        # Send count number of items
        for i in range(1, count + 1):
            # Check if we should stop
            if context.should_stop():
                break

            # Send the squared number
            context.output_queue.put(i * i)
            # Minimal delay to avoid flooding the queue while keeping tests fast
            time.sleep(0.001)

        # Signal end of stream
        context.output_queue.put(None)


@dataclass
class ThreadedMultiplier(ParallelizeTransformElement):
    """A transform element that multiplies input by a factor using threading."""

    _use_threading_override: bool = True
    multiplier: int = 2
    at_eos: bool = False

    def __post_init__(self):
        super().__post_init__()
        self.frame_list = []

    def pull(self, pad, frame):
        if self.in_queue is not None:
            self.in_queue.put(frame)
        if frame.EOS and not self.terminated.is_set():
            self.at_eos = True
            self.frame_list = self.sub_process_shutdown(10)

    def worker_process(self, context: WorkerContext, multiplier: int) -> None:
        try:
            frame = context.input_queue.get(timeout=0.1)
            if frame:
                if not frame.EOS:
                    # Modify the frame data
                    frame.data = frame.data * multiplier
                context.output_queue.put(frame)
        except queue.Empty:
            pass

    def new(self, pad):
        if not self.frame_list:
            return self.out_queue.get()
        return self.frame_list.pop(0)


@dataclass
class ProcessedSquarer(ParallelizeTransformElement):
    """A transform element that squares input using multiprocessing."""

    _use_threading_override: bool = False
    at_eos: bool = False

    def __post_init__(self):
        super().__post_init__()
        self.frame_list = []

    def pull(self, pad, frame):
        if self.in_queue is not None:
            self.in_queue.put(frame)
        if frame.EOS and not self.terminated.is_set():
            self.at_eos = True
            self.frame_list = self.sub_process_shutdown(10)

    def worker_process(self, context):
        try:
            frame = context.input_queue.get(timeout=0.1)
            if frame:
                if not frame.EOS:
                    # Square the data
                    frame.data = frame.data**2
                context.output_queue.put(frame)
        except queue.Empty:
            pass

    def new(self, pad):
        if not self.frame_list:
            return self.out_queue.get()
        return self.frame_list.pop(0)


@dataclass
class ResultCollector(ParallelizeSinkElement):
    """A sink element that collects results for testing."""

    _use_threading_override: bool = True
    at_eos: bool = False

    def __post_init__(self):
        super().__post_init__()
        self.results = {}
        for pad_name in self.sink_pad_names:
            self.results[pad_name] = []
        self.eos_count = 0
        self.expected_eos_count = len(self.sink_pad_names)

    def pull(self, pad, frame):
        if frame.EOS:
            self.mark_eos(pad)
            self.eos_count += 1
            # Only consider at_eos when all pads have received EOS
            if self.eos_count >= self.expected_eos_count:
                self.at_eos = True
                # Shutdown the worker when all pads have received EOS
                if not self.terminated.is_set():
                    self.sub_process_shutdown(
                        30
                    )  # Increased timeout for CI environments

        # Only send to queue if it exists
        if self.in_queue is not None:
            self.in_queue.put((pad.name, frame))

    def worker_process(self, context):
        try:
            data = context.input_queue.get(timeout=0.1)
            if data:
                pad_name, frame = data
                if not frame.EOS:
                    # Store the frame data in the results
                    context.output_queue.put((pad_name, frame.data))
        except queue.Empty:
            pass

    def get_results(self):
        """Get the collected results."""
        # Only read from queue if it exists
        if self.out_queue is not None:
            try:
                while True:
                    pad_name, value = self.out_queue.get_nowait()
                    self.results[pad_name].append(value)
            except (Empty, AttributeError):
                pass  # Queue is empty or was already closed
        return self.results


#
# Basic subprocess tests
#
def test_subprocess():
    """Test basic subprocess functionality with a simple pipeline."""
    shared_data = bytearray(
        "Here is a string that will be shared between processes", "utf-8"
    )
    Parallelize.to_shm("shared_data", shared_data)

    source = MySourceClass(source_pad_names=("event",))
    transform1 = MyTransformClass(
        sink_pad_names=("event",), source_pad_names=("samples1",)
    )
    transform2 = MyTransformClass(
        sink_pad_names=("event",), source_pad_names=("samples2",)
    )
    sink = MySinkClass(sink_pad_names=("samples1", "samples2"))

    pipeline = Pipeline()
    pipeline.insert(
        source,
        transform1,
        transform2,
        sink,
        link_map={
            sink.snks["samples1"]: transform1.srcs["samples1"],
            sink.snks["samples2"]: transform2.srcs["samples2"],
            transform1.snks["event"]: source.srcs["event"],
            transform2.snks["event"]: source.srcs["event"],
        },
    )

    # Use automatic parallelization detection
    pipeline.run()


def test_subprocess_exit_kill():
    """Test that __exit__ method's kill branch is executed."""
    # Setup instance list with a stub worker
    kill_called = [False]
    cancel_join_called = [False, False]  # For in_queue and out_queue

    class StubWorker:
        def __init__(self):
            self.alive = True

        def is_alive(self):
            return self.alive

        def start(self):
            # Stub for starting the worker
            pass

        def join(self, timeout):
            # Join doesn't terminate the worker
            pass

        def kill(self):
            kill_called[0] = True
            self.alive = False

    class StubQueue:
        def cancel_join_thread(self):
            # Keep track of cancel_join_thread calls
            if self.is_in_queue:
                cancel_join_called[0] = True
            else:
                cancel_join_called[1] = True

        def __init__(self, is_in_queue=True):
            self.is_in_queue = is_in_queue

    class StubInstance:
        def __init__(self):
            self.worker = StubWorker()
            self.in_queue = StubQueue(is_in_queue=True)
            self.out_queue = StubQueue(is_in_queue=False)
            self.worker_stop = multiprocessing.Event()

    # Use monkeypatch to temporarily replace instance_list
    original_instances = Parallelize.instance_list
    Parallelize.instance_list = [StubInstance()]

    try:
        # Create and use a minimal test pipeline
        test_pipeline = Pipeline()
        parallelize = Parallelize(test_pipeline)

        # The cleanup happens in __exit__
        with parallelize:
            pass

        # Verify that kill was called
        assert kill_called[0], "kill() method was not called in __exit__"
        # Verify that cancel_join_thread was called on both queues
        assert cancel_join_called[0], "cancel_join_thread not called on in_queue"
        assert cancel_join_called[1], "cancel_join_thread not called on out_queue"
    finally:
        # Restore original state
        Parallelize.instance_list = original_instances


def test_subprocess_run_exception():
    """Test that the run method properly handles exceptions in the pipeline."""

    class TestPipeline:
        def run(self, auto_parallelize=True):
            raise ValueError("Test exception")

    # Create a custom process class with a kill method that we can track
    kill_called = [False]

    class MockProcess:
        def is_alive(self):
            return True

        def join(self, timeout):
            # Simulate that join doesn't terminate the process
            pass

        def kill(self):
            kill_called[0] = True

    class MockLegacyProcess:
        def is_alive(self):
            return True

        def join(self, timeout):
            pass

        # No kill method

    class MockInstance:
        def __init__(self):
            # Use only the new naming convention
            self.worker = MockProcess()
            self.in_queue = multiprocessing.Queue(maxsize=1)
            self.out_queue = multiprocessing.Queue(maxsize=1)
            self.worker_stop = multiprocessing.Event()

    class MockLegacyInstance:
        def __init__(self):
            # Make this compatible with the new naming convention
            self.worker = MockLegacyProcess()
            self.in_queue = None  # Test the in_queue is None path
            self.out_queue = None  # Test the out_queue is None path
            self.worker_stop = multiprocessing.Event()

    # Add to the instance list so it gets cleaned up
    mock_instance = MockInstance()
    mock_legacy_instance = MockLegacyInstance()
    Parallelize.instance_list.extend([mock_instance, mock_legacy_instance])

    # Now run the test
    parallelize = Parallelize(TestPipeline())
    with pytest.raises(ValueError):
        parallelize.run()

    # Verify cleanup
    assert mock_instance.worker_stop.is_set(), "Worker stop event should be set"
    assert mock_legacy_instance.worker_stop.is_set(), "Worker stop event should be set"
    assert kill_called[0], "kill() method was not called"


#
# Test low-level subprocess wrapper components
#
def test_subprocess_wrapper():
    """Test the basic operation of _worker_wrapper_function."""
    terminated = multiprocessing.Event()
    shutdown = multiprocessing.Event()
    stop = multiprocessing.Event()
    shutdown.set()
    stop.set()
    inq = multiprocessing.Queue(maxsize=1)
    outq = multiprocessing.Queue(maxsize=1)

    class TestElement(_ParallelizeBase):
        def worker_process(self, context):
            pass

    _worker_wrapper_function(
        terminated,
        TestElement,
        "worker_process",
        worker_shutdown=shutdown,
        worker_stop=stop,
        inq=inq,
        outq=outq,
    )


def test_subprocess_wrapper_with_exception():
    """Test _worker_wrapper_function with a function that raises an exception."""
    terminated = multiprocessing.Event()
    shutdown = multiprocessing.Event()
    stop = multiprocessing.Event()
    inq = multiprocessing.Queue(maxsize=1)
    outq = multiprocessing.Queue(maxsize=1)

    class TestElement(_ParallelizeBase):
        def worker_process(self, context):
            raise RuntimeError("nope")

    _worker_wrapper_function(
        terminated,
        TestElement,
        "worker_process",
        worker_shutdown=shutdown,
        worker_stop=stop,
        inq=inq,
        outq=outq,
    )

    # Terminated should be set even with an exception
    assert terminated.is_set(), "terminated event was not set after exception"


def test_subprocess_wrapper_with_threading():
    """Test _worker_wrapper_function with threading."""
    terminated = threading.Event()
    shutdown = threading.Event()
    stop = threading.Event()
    shutdown.set()
    inq = multiprocessing.Queue(maxsize=1)
    inq.put(None)
    outq = multiprocessing.Queue(maxsize=1)
    outq.put(None)

    class TestElement(_ParallelizeBase):
        def worker_process(self, context):
            raise ValueError("nope")

    thread = threading.Thread(
        target=_worker_wrapper_function,
        args=(terminated, TestElement, "worker_process"),
        kwargs={
            "worker_shutdown": shutdown,
            "worker_stop": stop,
            "inq": inq,
            "outq": outq,
        },
    )
    thread.start()
    time.sleep(1)
    stop.set()
    shutdown.set()
    thread.join()


def test_subprocess_keyboard_interrupt():
    """Test that KeyboardInterrupt is properly caught and handled."""
    # Set up events and queues
    terminated = threading.Event()
    shutdown = threading.Event()
    stop = threading.Event()
    inq = multiprocessing.Queue(maxsize=2)
    outq = multiprocessing.Queue(maxsize=2)

    # Flag to track iterations
    iteration_count = 0
    keyboard_interrupt_raised = False
    completed_after_interrupt = False

    class TestElement(_ParallelizeBase):
        def worker_process(self, context):
            nonlocal iteration_count, keyboard_interrupt_raised
            nonlocal completed_after_interrupt

            # First call: raise KeyboardInterrupt
            if iteration_count == 0:
                iteration_count += 1
                keyboard_interrupt_raised = True
                raise KeyboardInterrupt("Test interrupt")

            # Second call: mark that we continued after the interrupt
            elif iteration_count == 1:
                iteration_count += 1
                completed_after_interrupt = True
                # Signal to stop now
                stop.set()

    # Run the wrapper in a thread
    def run_wrapper():
        _worker_wrapper_function(
            terminated,
            TestElement,
            "worker_process",
            worker_shutdown=shutdown,
            worker_stop=stop,
            inq=inq,
            outq=outq,
        )

    thread = threading.Thread(target=run_wrapper)
    thread.daemon = True
    thread.start()

    # Wait for the thread to complete
    thread.join(timeout=1)

    # Verify that we continued after the KeyboardInterrupt
    assert keyboard_interrupt_raised, "KeyboardInterrupt was not raised"
    assert (
        completed_after_interrupt
    ), "Execution did not continue after KeyboardInterrupt"
    assert terminated.is_set(), "The terminated event should be set"


def test_subprocess_drain_queue():
    """Test the queue draining logic in _sub_process_wrapper during orderly shutdown."""
    # Set up events and queues - use threading.Event for consistent behavior
    terminated = threading.Event()
    worker_shutdown = threading.Event()
    worker_stop = threading.Event()

    # Set shutdown but not stop - this is the key condition for drain logic
    worker_shutdown.set()

    # Set up queue with items to process - fewer items for faster tests
    inq = multiprocessing.Queue(maxsize=3)
    for i in range(3):
        inq.put(Frame(data=f"Test Item {i}", EOS=False))
    outq = multiprocessing.Queue(maxsize=3)

    # Track calls to func
    call_count = 0

    class TestElement(_ParallelizeBase):
        def worker_process(self, context):
            nonlocal call_count
            try:
                frame = context.input_queue.get(block=False)
                if frame:
                    call_count += 1
                    # Simulate processing by printing
                    print(f"Processing item: {frame.data}")
                    # Explicitly set the terminated event at the end
                    terminated.set()
            except queue.Empty:
                pass

    # Use a thread so we can set process_stop after a delay
    def run_wrapper():
        _worker_wrapper_function(
            terminated,
            TestElement,
            "worker_process",
            worker_shutdown=worker_shutdown,
            worker_stop=worker_stop,
            inq=inq,
            outq=outq,
        )

    # Start thread
    thread = threading.Thread(target=run_wrapper)
    thread.daemon = True
    thread.start()

    # Let it run for a bit to process the queue
    time.sleep(0.3)  # Slightly longer delay to ensure processing completes

    # Now set stop to allow the thread to exit
    worker_stop.set()
    thread.join(timeout=1)

    # Verify items were processed
    assert (
        call_count >= 3
    ), f"Expected at least 3 calls to process items, got {call_count}"
    assert terminated.is_set(), "The terminated event should be set"


def test_subprocess_internal_not_implemented():
    """Test that _ParallelizeBase.worker_process raises NotImplementedError.

    This confirms correct base class behavior.
    """
    base = _ParallelizeBase()
    with pytest.raises(NotImplementedError):
        base.worker_process(None)


def test_subprocess_internal_runtime_error():
    """Test for RuntimeError from internal when terminated before EOS.

    Verifies correct error handling when worker terminates prematurely.
    """

    class TestParallelizeElement(_ParallelizeBase):
        def __init__(self):
            self.terminated = multiprocessing.Event()
            self.terminated.set()  # Set terminated
            self.at_eos = False  # But not at_eos
            self._retrieved_worker_exception = None

    element = TestParallelizeElement()
    with pytest.raises(RuntimeError):
        element.internal()


def test_subprocess_shutdown_timeout():
    """Test that sub_process_shutdown raises RuntimeError on timeout."""

    class TestParallelizeElement(_ParallelizeBase):
        def __init__(self):
            self.worker_shutdown = multiprocessing.Event()
            self.worker_stop = multiprocessing.Event()
            self.terminated = multiprocessing.Event()
            self.out_queue = multiprocessing.Queue()
            self.in_queue = multiprocessing.Queue()

    element = TestParallelizeElement()
    # Set a very small timeout to trigger the timeout error
    with pytest.raises(RuntimeError):
        element.sub_process_shutdown(timeout=0.001)


def test_subprocess_to_shm_duplicate():
    """Test that attempting to create duplicate shared memory raises FileExistsError."""
    # First clear any existing shared memory with this name
    try:
        Parallelize.shm_list = []
        multiprocessing.shared_memory.SharedMemory(name="test_duplicate").unlink()
    except FileNotFoundError:
        # This is fine - means the memory segment doesn't exist yet
        pass

    # Create the first shared memory instance
    test_data = bytearray("Test data for shared memory duplicate test", "utf-8")
    Parallelize.to_shm("test_duplicate", test_data)

    # Now try to create another with the same name, which should fail
    with pytest.raises(FileExistsError):
        Parallelize.to_shm("test_duplicate", test_data)

    # Clean up
    for d in Parallelize.shm_list:
        try:
            multiprocessing.shared_memory.SharedMemory(name=d["name"]).unlink()
        except FileNotFoundError:
            pass
    Parallelize.shm_list = []


#
# Tests for concurrency modes
#
def test_threading_mode():
    """Test the threading concurrency mode."""
    # Set global default to use threading for this test
    Parallelize.use_threading_default = True

    # Create elements with minimal count (just 1) to ensure quick test completion
    source = NumberSource(count=1, source_pad_names=("numbers",))
    transform = ThreadedMultiplier(
        sink_pad_names=("in",),
        source_pad_names=("out",),
        multiplier=3,
    )
    collector = ResultCollector(sink_pad_names=("original", "transformed"))

    # Create and set up the pipeline
    pipeline = Pipeline()
    pipeline.insert(
        source,
        transform,
        collector,
        link_map={
            transform.snks["in"]: source.srcs["numbers"],
            collector.snks["original"]: source.srcs["numbers"],
            collector.snks["transformed"]: transform.srcs["out"],
        },
    )

    # Run with default (which is True for this test)
    with Parallelize(pipeline) as parallelize:
        parallelize.run()

    # Get the results
    results = collector.get_results()

    # Just verify we got some results and basic correctness
    # No need to check lengths or iterate through multiple values
    if results.get("original"):
        assert results["original"][0] == 1
    if results.get("transformed"):
        assert results["transformed"][0] == 3


def test_mixed_concurrency():
    """Test mixing threading and multiprocessing in the same pipeline."""
    # Default to process mode for this test
    Parallelize.use_threading_default = False

    # Create elements with explicit concurrency modes - use just 1 item for speed
    source = NumberSource(count=1, source_pad_names=("numbers",))

    thread_transform = ThreadedMultiplier(
        sink_pad_names=("in",),
        source_pad_names=("out",),
        multiplier=2,
    )

    process_transform = ProcessedSquarer(
        sink_pad_names=("in",),
        source_pad_names=("out",),
    )

    # Use thread mode for collector
    collector = ResultCollector(sink_pad_names=("original", "doubled", "squared"))

    # Create and set up the pipeline
    pipeline = Pipeline()
    pipeline.insert(
        source,
        thread_transform,
        process_transform,
        collector,
        link_map={
            thread_transform.snks["in"]: source.srcs["numbers"],
            process_transform.snks["in"]: thread_transform.srcs["out"],
            collector.snks["original"]: source.srcs["numbers"],
            collector.snks["doubled"]: thread_transform.srcs["out"],
            collector.snks["squared"]: process_transform.srcs["out"],
        },
    )

    # Run the pipeline with default mode
    with Parallelize(pipeline) as parallelize:
        parallelize.run()

    # Get the results
    results = collector.get_results()

    # Just check for the presence of expected values using the minimal validation needed
    if results.get("original"):
        assert results["original"][0] == 1
    if results.get("doubled"):
        assert results["doubled"][0] == 2
    if results.get("squared"):
        assert results["squared"][0] == 4  # 2²=4


def test_subprocess_source_process():
    """Test a simple process source element implementation.

    This is a simplified test that only checks that the source element correctly
    implements the abstract new method.
    """
    # Create a source element directly without a pipeline
    source = SimpleProcessSource(count=1, source_pad_names=("output",))

    # Verify that the source element has a new method
    assert hasattr(source, "new")

    # Check that the new method is callable
    # (Using source.new directly instead of getattr to avoid linter warnings)
    assert callable(source.new)

    # Calling the new method should not raise NotImplementedError
    try:
        frame = source.new(source.source_pads[0])
        # Frame might be empty since the worker isn't running, which is fine
        assert isinstance(frame, Frame)
    except NotImplementedError:
        # Raise the assertion error directly instead of using assert False
        raise AssertionError("new method should be implemented, not abstract")


def test_complete_subprocess_pipeline():
    """Test that a source element with the abstract new method
    can be instantiated and used in a pipeline.

    This is a minimal test that just verifies the source element can be
    created and that it implements the abstract new method correctly.
    """
    # Create a simple source element
    source = SimpleProcessSource(count=1, source_pad_names=("numbers",))

    # Verify that the source element has a new method
    assert hasattr(source, "new")

    # Check that the new method is callable
    assert callable(source.new)

    # Verify that the pad_eos_sent dictionary has entries and none are True
    assert len(source.pad_eos_sent) > 0
    assert all(not v for v in source.pad_eos_sent.values())


def test_worker_context_comprehensive():
    """Comprehensive test for WorkerContext functionality."""

    # Test WorkerContext edge cases with None values
    context = WorkerContext()
    assert context.input_queue is None
    assert context.output_queue is None
    assert context.stop_event is None
    assert not context.should_stop() and not context.should_shutdown()

    # Test with actual events and queues
    import threading

    stop_event = threading.Event()
    shutdown_event = threading.Event()
    input_q = queue.Queue()
    output_q = queue.Queue()

    context = WorkerContext(
        input_queue=input_q,
        output_queue=output_q,
        worker_stop=stop_event,
        worker_shutdown=shutdown_event,
    )

    # Test queue access
    assert context.input_queue is input_q
    assert context.output_queue is output_q

    # Test event handling
    stop_event.set()
    assert context.should_stop() is True

    shutdown_event.set()
    assert context.should_shutdown() is True


def test_pipeline_parallelization_detection():
    """Test automatic parallelization detection for pipelines."""

    # Empty pipeline should not need parallelization
    empty_pipeline = Pipeline()
    assert Parallelize.needs_parallelization(empty_pipeline) is False

    # Pipeline with regular elements should not need parallelization
    class RegularSource(SourceElement):
        def new(self, pad):
            return Frame(data=None, EOS=True)

    regular_pipeline = Pipeline()
    regular_pipeline.insert(RegularSource(source_pad_names=("out",)))
    assert Parallelize.needs_parallelization(regular_pipeline) is False


def test_parameter_extraction_edge_cases():
    # Create a minimal class that has _extract_worker_parameters but no worker_process
    class ElementWithoutWorkerProcess:
        def _extract_worker_parameters(self):
            """Extract parameters for worker_process method from instance attributes."""
            if not hasattr(self, "worker_process"):
                return {}
            # ... rest would be unreachable in this test

    element_no_worker = ElementWithoutWorkerProcess()
    extracted = element_no_worker._extract_worker_parameters()
    assert extracted == {}, "Should return empty dict when no worker_process method"

    # Parameter uses default value from method signature
    @dataclass
    class TestElementWithMethodDefaults(ParallelizeTransformElement):
        existing_param: int = 42

        def new(self, pad):
            return Frame()

        def pull(self, pad, frame):
            pass

        @staticmethod
        def worker_process(
            context: WorkerContext,
            existing_param: int,
            method_default_param: str = "from_method",
        ) -> None:
            pass

    # Test element where parameter has default in method but not as instance attribute
    element = TestElementWithMethodDefaults(
        sink_pad_names=("input",), source_pad_names=("output",)
    )

    # This should extract parameters and use method defaults where needed (line 446)
    extracted = element._extract_worker_parameters()

    assert "existing_param" in extracted
    assert "method_default_param" in extracted
    assert extracted["existing_param"] == 42  # From instance
    assert extracted["method_default_param"] == "from_method"


def test_worker_exception_storage_threading():
    """Test worker exception storage in threading mode."""

    @dataclass
    class ExceptionSource(ParallelizeSourceElement):
        """A source that raises an exception in worker_process."""

        _use_threading_override: bool = True  # Force threading mode

        def __post_init__(self):
            super().__post_init__()

        def new(self, pad):
            # Return empty frame - worker handles the logic
            return Frame(data=None)

        def worker_process(self, context: WorkerContext) -> None:
            """Worker that raises a test exception."""
            raise ValueError("Test exception from worker")

    # Create the source element
    source = ExceptionSource(source_pad_names=("test",))

    # Start the worker thread and let it fail
    source.worker.start()
    source.worker.join(timeout=2)

    # Verify that the worker terminated
    assert source.terminated.is_set()

    # Get the exception and verify it
    stored_exception = source.get_worker_exception()
    assert stored_exception is not None
    assert isinstance(stored_exception, ValueError)
    assert str(stored_exception) == "Test exception from worker"


# Module-level class for test_worker_exception_storage_multiprocessing
# Must be at module level to be pickleable by multiprocessing
@dataclass
class ExceptionSourceForTest(ParallelizeSourceElement):
    """A source that raises an exception in worker_process."""

    _use_threading_override: bool = False

    def __post_init__(self):
        super().__post_init__()

    def new(self, pad):
        return Frame(data=None)

    def worker_process(self, context: WorkerContext) -> None:
        """Worker that raises a test exception."""
        raise ValueError("Test exception from worker")


def test_worker_exception_storage_multiprocessing():
    """Test worker exception storage in multiprocessing mode."""
    # Create the source element
    source = ExceptionSourceForTest(source_pad_names=("test",))

    # Start the worker process and let it fail
    source.worker.start()
    source.worker.join(timeout=2)

    # Verify that the worker terminated
    assert source.terminated.is_set()

    # Get the exception and verify it
    stored_exception = source.get_worker_exception()
    assert stored_exception is not None
    assert isinstance(stored_exception, ValueError)
    assert str(stored_exception) == "Test exception from worker"


def test_get_worker_exception_threading():
    """Test get_worker_exception method in threading mode."""

    @dataclass
    class FailingThreadedSource(ParallelizeSourceElement):
        """Source that fails in worker with a specific error."""

        _use_threading_override: bool = True

        def __post_init__(self):
            super().__post_init__()

        def new(self, pad):
            return Frame(data=None)

        def worker_process(self, context: WorkerContext) -> None:
            raise AssertionError("Custom assertion error from threaded worker")

    # Create the element
    source = FailingThreadedSource(source_pad_names=("test",))

    # Start the worker thread and let it fail
    source.worker.start()
    source.worker.join(timeout=2)

    # Verify that the worker terminated and we can get the exception
    assert source.terminated.is_set()

    # Test that get_worker_exception returns the stored exception
    worker_exc = source.get_worker_exception()
    assert worker_exc is not None
    assert isinstance(worker_exc, AssertionError)
    assert str(worker_exc) == "Custom assertion error from threaded worker"

    # Test caching: second call should return the same cached exception
    worker_exc2 = source.get_worker_exception()
    assert worker_exc2 is worker_exc  # Same object, proving it's cached


# Module-level class for test_get_worker_exception_multiprocessing
# Must be at module level to be pickleable by multiprocessing
@dataclass
class FailingProcessSourceForTest(ParallelizeSourceElement):
    """Source that fails in worker with a specific error."""

    _use_threading_override: bool = False

    def __post_init__(self):
        super().__post_init__()

    def new(self, pad):
        return Frame(data=None)

    def worker_process(self, context: WorkerContext) -> None:
        raise ValueError("Custom value error from process worker")


def test_get_worker_exception_multiprocessing():
    """Test get_worker_exception method in multiprocessing mode."""
    # Create the element
    source = FailingProcessSourceForTest(source_pad_names=("test",))

    # Start the worker process and let it fail
    source.worker.start()
    source.worker.join(timeout=2)

    # Verify that the worker terminated and we can get the exception
    assert source.terminated.is_set()

    # Test that get_worker_exception returns the stored exception
    worker_exc = source.get_worker_exception()
    assert worker_exc is not None
    assert isinstance(worker_exc, ValueError)
    assert str(worker_exc) == "Custom value error from process worker"

    # Test caching: second call should return the same cached exception
    worker_exc2 = source.get_worker_exception()
    assert worker_exc2 is worker_exc  # Same object, proving it's cached


def test_get_worker_exception_pipeline():
    """Test that worker exceptions are properly chained through the pipeline."""

    @dataclass
    class ExceptionSource(ParallelizeSourceElement):
        """A source that raises an exception in worker_process."""

        _use_threading_override: bool = True  # Force threading mode

        def __post_init__(self):
            super().__post_init__()

        def new(self, pad):
            return Frame(data=None)

        def worker_process(self, context: WorkerContext) -> None:
            """Worker that raises a test exception."""
            raise ValueError("Test exception from worker")

    pipeline = Pipeline()
    pipeline.insert(
        ExceptionSource(
            name="src1",
            source_pad_names=("H1",),
        ),
        NullSink(
            name="snk1",
            sink_pad_names=("H1",),
        ),
        link_map={
            "snk1:snk:H1": "src1:src:H1",
        },
    )

    # The key test: pipeline should raise RuntimeError but preserve the
    # original ValueError in the chain
    with pytest.raises(RuntimeError, match="worker stopped before EOS") as exc_info:
        pipeline.run()

    # Verify the original ValueError is in the exception chain
    current_exc = exc_info.value
    found_value_error = False
    while current_exc is not None:
        if isinstance(current_exc, ValueError) and "Test exception from worker" in str(
            current_exc
        ):
            found_value_error = True
            break
        current_exc = current_exc.__cause__

    assert (
        found_value_error
    ), "Original ValueError should be preserved in exception chain"


if __name__ == "__main__":
    test_subprocess()
