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

import networkx
import pytest

from mergekitty.common import ImmutableMap, ModelReference
from mergekitty.executor import ParallelismExecutor, SingleThreadedExecutor
from mergekitty.io.tasks import LoadTensor, LoaderCache
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"
        )

    def test_load_tensor_scheduling_groups_by_shard(self, executor_cls, monkeypatch):
        model = ModelReference.parse("fake-model")
        tensor_paths = {
            "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
            "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
            "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
            "model.layers.10.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
            "model.layers.11.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
            "model.layers.12.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
        }
        fake_loader = SimpleNamespace(
            index=SimpleNamespace(tensor_paths=tensor_paths),
        )
        monkeypatch.setattr(LoaderCache, "get", lambda self, _model: fake_loader)

        tasks = [LoadTensor(model=model, tensor=name) for name in tensor_paths]
        schedule = executor_cls(tasks)._make_schedule(tasks)

        shard_switches = 0
        previous_shard = None
        for task in schedule:
            shard = tensor_paths[task.tensor]
            if previous_shard is not None and shard != previous_shard:
                shard_switches += 1
            previous_shard = shard

        assert shard_switches <= 1, (
            "LoadTensor scheduling should preserve shard locality"
        )


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()
