#!/usr/bin/env python3

from __future__ import annotations

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

from sgn.sources import SignalEOS
from sgn.subprocess import (
    Parallelize,
    _ParallelizeBase,
    ParallelizeTransformElement,
    ParallelizeSinkElement,
    ParallelizeSourceElement,
)
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(10)

    @staticmethod
    def sub_process_internal(**kwargs):
        kwargs["outq"].put(None)
        try:
            kwargs["inq"].get(timeout=0.1)
        except 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)

    @staticmethod
    def sub_process_internal(**kwargs):
        # access some shared memory - there is only one
        # Just access it to verify it exists, but don't need to use it
        _ = kwargs["shm_list"][0]["shm"]
        try:
            frame = kwargs["inq"].get(timeout=0.1)
            kwargs["outq"].put(frame)
        except 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 = 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}
        # Store count in worker arguments
        self.worker_argdict = {"count": self.count}
        # 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)

    @staticmethod
    def sub_process_internal(**kwargs):
        """Generate sequential numbers and send to the main thread."""
        outq = kwargs["outq"]
        worker_stop = kwargs["worker_stop"]
        worker_argdict = kwargs.get("worker_argdict", {})
        count = worker_argdict.get("count", 3)

        # Send count number of items
        for i in range(1, count + 1):
            # Check if we should stop
            if worker_stop.is_set():
                break

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

        # Signal end of stream
        outq.put(None)


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

    _use_threading_override = 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}
        # Store count in worker arguments
        self.worker_argdict = {"count": self.count}
        # 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)

    @staticmethod
    def sub_process_internal(**kwargs):
        """Generate squared numbers and send to the main process."""
        outq = kwargs["outq"]
        worker_stop = kwargs["worker_stop"]
        worker_argdict = kwargs.get("worker_argdict", {})
        count = worker_argdict.get("count", 3)

        # Send count number of items
        for i in range(1, count + 1):
            # Check if we should stop
            if worker_stop.is_set():
                break

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

        # Signal end of stream
        outq.put(None)


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

    _use_threading_override = 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)

    @staticmethod
    def sub_process_internal(**kwargs):
        inq, outq = kwargs["inq"], kwargs["outq"]
        worker_argdict = kwargs.get("worker_argdict", {})
        multiplier = worker_argdict.get("multiplier", 2)

        try:
            frame = inq.get(timeout=0.1)
            if not frame.EOS:
                # Modify the frame data
                frame.data = frame.data * multiplier
            outq.put(frame)
        except 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 = 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)

    @staticmethod
    def sub_process_internal(**kwargs):
        inq, outq = kwargs["inq"], kwargs["outq"]
        try:
            frame = inq.get(timeout=0.1)
            if not frame.EOS:
                # Square the data
                frame.data = frame.data**2
            outq.put(frame)
        except 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 = 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(10)

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

    @staticmethod
    def sub_process_internal(**kwargs):
        inq, outq = kwargs["inq"], kwargs["outq"]
        # Worker stop event is available but not needed for this test
        _ = kwargs.get("worker_stop")

        try:
            pad_name, frame = inq.get(timeout=0.1)
            if not frame.EOS and outq is not None:
                # Store the frame data in the results
                outq.put((pad_name, frame.data))
        except 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"],
        },
    )

    with Parallelize(pipeline) as parallelize:
        parallelize.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_normal_completion():
    """Test that worker_stop events are set when pipeline completes normally."""

    # Test pipeline that completes normally
    class TestPipeline:
        def run(self):
            # Just return successfully
            pass

    # Create test instance
    class TestInstance:
        def __init__(self):
            self.worker_stop = multiprocessing.Event()

    # Save original and create test instances
    original_instances = Parallelize.instance_list.copy()
    Parallelize.instance_list = []

    try:
        # Add our test instance
        instance = TestInstance()
        Parallelize.instance_list.append(instance)

        # Create a Parallelize and run it normally
        parallelize = Parallelize(TestPipeline())
        parallelize.run()

        # Verify the stop event was set
        assert (
            instance.worker_stop.is_set()
        ), "worker_stop event was not set on normal completion"
    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):
            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(RuntimeError):
        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 _sub_process_wrapper."""
    terminated = multiprocessing.Event()
    shutdown = multiprocessing.Event()
    stop = multiprocessing.Event()
    shutdown.set()
    stop.set()
    inq = multiprocessing.Queue(maxsize=1)
    outq = multiprocessing.Queue(maxsize=1)

    def func(**kwargs):
        pass

    _ParallelizeBase._sub_process_wrapper(
        func,
        terminated,
        worker_shutdown=shutdown,
        worker_stop=stop,
        inq=inq,
        outq=outq,
    )


def test_subprocess_wrapper_with_exception():
    """Test _sub_process_wrapper 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)

    def func(**kwargs):
        raise RuntimeError("nope")

    _ParallelizeBase._sub_process_wrapper(
        func,
        terminated,
        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 _sub_process_wrapper 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)

    def func(**kwargs):
        raise ValueError("nope")

    thread = threading.Thread(
        target=_ParallelizeBase._sub_process_wrapper,
        args=(func, terminated),
        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

    def test_func(**kwargs):
        nonlocal iteration_count, keyboard_interrupt_raised, 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():
        _ParallelizeBase._sub_process_wrapper(
            test_func,
            terminated,
            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

    # This function will be called to process each item from the queue
    def test_func(**kwargs):
        nonlocal call_count
        q = kwargs["inq"]
        try:
            item = q.get(block=False)
            call_count += 1
            # Simulate processing by printing
            print(f"Processing item: {item.data}")
            # Explicitly set the terminated event at the end
            terminated.set()
        except Empty:
            pass

    # Use a thread so we can set process_stop after a delay
    def run_wrapper():
        _ParallelizeBase._sub_process_wrapper(
            test_func,
            terminated,
            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.sub_process_internal raises NotImplementedError.

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


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

    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",),
        worker_argdict={"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",),
        worker_argdict={"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())


if __name__ == "__main__":
    test_subprocess()
