from slurm.callbacks import BaseCallback
from slurm.decorators import task
from slurm.packaging.base import PackagingStrategy
from cluster_factory import make_test_cluster  # type: ignore
from local_backend import LocalBackend  # type: ignore


class NoopStrategy(PackagingStrategy):
    def prepare(self, task, cluster):
        return {"status": "ok"}

    def generate_setup_commands(self, task, job_id=None, job_dir=None):
        return []

    def generate_cleanup_commands(self, task, job_id=None, job_dir=None):
        return []


@task(job_name="sum", time="00:01:00", mem="1G", ntasks=1, nodes=1)
def add_one(x: int) -> int:
    return x + 1


def test_cluster_submit_with_local_backend(monkeypatch, tmp_path):
    captured: dict[str, object] = {}

    class CaptureCallback(BaseCallback):
        def on_begin_submit_job_ctx(self, ctx):
            captured["begin"] = ctx

        def on_end_submit_job_ctx(self, ctx):
            captured["end"] = ctx

    cluster = make_test_cluster(
        backend=LocalBackend(job_base_dir=str(tmp_path)),
        callbacks=[CaptureCallback()],
    )

    # Submit and run locally
    job = cluster.submit(add_one, packaging="none")(41)

    # Job should be immediately completed in local backend
    assert job.is_completed()
    assert job.is_successful()
    assert job.get_result() == 42

    begin_ctx = captured.get("begin")
    end_ctx = captured.get("end")

    assert begin_ctx is not None, "Submit begin context not captured"
    assert end_ctx is not None, "Submit end context not captured"

    assert getattr(begin_ctx, "backend_type", None) == "LocalBackend"
    assert getattr(end_ctx, "backend_type", None) == "LocalBackend"
