# Partly Copyright (C) 2025 Arcee AI
# Partly Copyright (C) 2025-2026 Allura-org
import threading
from typing import Any, Dict, Optional

import networkx
import pytest

from mergekitty.common import ImmutableMap
from mergekitty.executor import ParallelismExecutor, SingleThreadedExecutor
from mergekitty.task import Task

EXECUTION_COUNTS: Dict[Task, int] = {}
BLOCKING_TASKS_STARTED: set[str] = set()
BLOCKING_TASKS_LOCK = threading.Lock()
BLOCKING_TASKS_READY = threading.Event()
BLOCKING_TASKS_RELEASE = threading.Event()


class DummyTask(Task):
    result: Any
    dependencies: ImmutableMap[str, Task]
    name: str = "DummyTask"
    grouplabel: Optional[str] = None
    execution_count: int = 0

    def arguments(self):
        return self.dependencies

    def group_label(self) -> Optional[str]:
        return self.grouplabel

    def execute(self, **kwargs):
        EXECUTION_COUNTS[self] = EXECUTION_COUNTS.get(self, 0) + 1
        return self.result


def create_mock_task(name, result=None, dependencies=None, group_label=None):
    if dependencies is None:
        dependencies = {}
    return DummyTask(
        result=result,
        dependencies=ImmutableMap(data=dependencies),
        name=name,
        grouplabel=group_label,
    )


@pytest.fixture(
    params=[SingleThreadedExecutor, ParallelismExecutor],
    ids=["single", "parallel"],
)
def executor_cls(request):
    return request.param


# Test cases for the Task implementation
class TestTaskClass:
    def test_task_execute(self):
        # Testing the execute method
        task = create_mock_task("task1", result=42)
        assert task.execute() == 42, "Task execution did not return expected result"

    def test_task_priority(self):
        task = create_mock_task("task1")
        assert task.priority() == 0, "Default priority should be 0"

    def test_task_group_label(self):
        task = create_mock_task("task1")
        assert task.group_label() is None, "Default group label should be None"


# Test cases for the Executor implementation
class TestExecutorClass:
    def test_executor_initialization(self, executor_cls):
        # Testing initialization with single task
        task = create_mock_task("task1")
        executor = executor_cls([task])
        assert executor.targets == [task], (
            "Executor did not initialize with correct targets"
        )

    def test_executor_empty_list(self, executor_cls):
        list(executor_cls([]).run())

    def test_executor_scheduling(self, executor_cls):
        # Testing scheduling with dependencies
        task1 = create_mock_task("task1", result=1)
        task2 = create_mock_task("task2", result=2, dependencies={"task1": task1})
        executor = executor_cls([task2])
        assert len(executor._make_schedule([task2])) == 2, (
            "Schedule should include two tasks"
        )

    def test_executor_dependency_building(self, executor_cls):
        # Testing dependency building
        task1 = create_mock_task("task1")
        task2 = create_mock_task("task2", dependencies={"task1": task1})
        executor = executor_cls([task2])
        dependencies = executor._build_dependencies([task2])
        assert task1 in dependencies[task2], "Task1 should be a dependency of Task2"

    def test_executor_run(self, executor_cls):
        # Testing execution through the run method
        task1 = create_mock_task("task1", result=10)
        task2 = create_mock_task("task2", result=20, dependencies={"task1": task1})
        executor = executor_cls([task2])
        results = list(executor.run())
        assert len(results) == 1 and results[0][1] == 20, (
            "Executor run did not yield correct results"
        )

    def test_executor_execute(self, executor_cls):
        # Testing execute method for side effects
        task1 = create_mock_task("task1", result=10)
        executor = executor_cls([task1])
        # No assert needed; we're ensuring no exceptions are raised and method completes
        executor.execute()

    def test_dependency_ordering(self, executor_cls):
        # Testing the order of task execution respects dependencies
        task1 = create_mock_task("task1", result=1)
        task2 = create_mock_task("task2", result=2, dependencies={"task1": task1})
        task3 = create_mock_task("task3", result=3, dependencies={"task2": task2})
        executor = executor_cls([task3])

        schedule = executor._make_schedule([task3])
        assert schedule.index(task1) < schedule.index(task2), (
            "Task1 should be scheduled before Task2"
        )
        assert schedule.index(task2) < schedule.index(task3), (
            "Task2 should be scheduled before Task3"
        )


class TestExecutorGroupLabel:
    def test_group_label_scheduling(self, executor_cls):
        # Create tasks with group labels and dependencies
        task1 = create_mock_task("task1", group_label="group1")
        task2 = create_mock_task(
            "task2", dependencies={"task1": task1}, group_label="group1"
        )
        task3 = create_mock_task("task3", group_label="group2")
        task4 = create_mock_task(
            "task4", dependencies={"task2": task2, "task3": task3}, group_label="group1"
        )

        # Initialize Executor with the tasks
        executor = executor_cls([task4])

        # Get the scheduled tasks
        schedule = executor._make_schedule([task4])

        # Check if tasks with the same group label are scheduled consecutively when possible
        group_labels_in_order = [
            task.group_label() for task in schedule if task.group_label()
        ]
        assert group_labels_in_order == [
            "group1",
            "group1",
            "group2",
            "group1",
        ], "Tasks with same group label are not scheduled consecutively"

    def test_group_label_with_dependencies(self, executor_cls):
        # Creating tasks with dependencies and group labels
        task1 = create_mock_task("task1", result=1, group_label="group1")
        task2 = create_mock_task(
            "task2", result=2, dependencies={"task1": task1}, group_label="group2"
        )
        task3 = create_mock_task(
            "task3", result=3, dependencies={"task2": task2}, group_label="group1"
        )

        executor = executor_cls([task3])
        schedule = executor._make_schedule([task3])
        scheduled_labels = [
            task.group_label() for task in schedule if task.group_label()
        ]

        # Check if task3 is scheduled after task1 and task2 due to dependency, even though it has the same group label as task1
        group1_indices = [
            i for i, label in enumerate(scheduled_labels) if label == "group1"
        ]
        group2_index = scheduled_labels.index("group2")

        assert group1_indices[-1] > group2_index, (
            "Task with the same group label but later dependency was not scheduled after different group label"
        )


class TestExecutorSingleExecution:
    def test_single_execution_per_task(self, executor_cls):
        EXECUTION_COUNTS.clear()

        shared_task = create_mock_task("shared_task", result=100)
        task1 = create_mock_task("task1", dependencies={"shared": shared_task})
        task2 = create_mock_task("task2", dependencies={"shared": shared_task})
        task3 = create_mock_task("task3", dependencies={"task1": task1, "task2": task2})

        executor_cls([task3]).execute()

        assert shared_task in EXECUTION_COUNTS, "Dependency not executed"
        assert EXECUTION_COUNTS[shared_task] == 1, (
            "Shared dependency should be executed exactly once"
        )


class CircularTask(Task):
    def arguments(self) -> Dict[str, Task]:
        return {"its_a_me": self}

    def execute(self, **_kwargs) -> Any:
        assert False, "Task with circular dependency executed"


class BlockingTask(Task):
    name: str

    def arguments(self) -> Dict[str, Task]:
        return {}

    def execute(self, **_kwargs) -> str:
        with BLOCKING_TASKS_LOCK:
            BLOCKING_TASKS_STARTED.add(self.name)
            if len(BLOCKING_TASKS_STARTED) >= 2:
                BLOCKING_TASKS_READY.set()

        assert BLOCKING_TASKS_RELEASE.wait(timeout=5), (
            "Timed out waiting for blocking tasks to be released"
        )
        return self.name


class TestExecutorCircularDependency:
    def test_circular_dependency(self, executor_cls):
        with pytest.raises(networkx.NetworkXUnfeasible):
            executor_cls([CircularTask()]).execute()


class TestParallelismExecutor:
    def test_executes_ready_tasks_in_parallel(self):
        with BLOCKING_TASKS_LOCK:
            BLOCKING_TASKS_STARTED.clear()
        BLOCKING_TASKS_READY.clear()
        BLOCKING_TASKS_RELEASE.clear()

        left = BlockingTask(name="left")
        right = BlockingTask(name="right")
        target = create_mock_task(
            "target",
            dependencies={"left": left, "right": right},
        )

        executor = ParallelismExecutor([target], max_workers=2)
        errors = []

        def _run_executor():
            try:
                executor.execute()
            except Exception as exc:  # pragma: no cover - assertion handles this
                errors.append(exc)
                BLOCKING_TASKS_RELEASE.set()

        thread = threading.Thread(target=_run_executor, daemon=True)
        thread.start()

        assert BLOCKING_TASKS_READY.wait(timeout=5), (
            "Parallel executor did not start both ready tasks concurrently"
        )
        with BLOCKING_TASKS_LOCK:
            assert BLOCKING_TASKS_STARTED == {"left", "right"}

        BLOCKING_TASKS_RELEASE.set()
        thread.join(timeout=5)

        assert not thread.is_alive(), "Parallel executor did not finish"
        assert not errors


if __name__ == "__main__":
    pytest.main()
