from __future__ import annotations

import unittest

from govoplan_core.core.dataflows import (
    CAPABILITY_DATAFLOW_RUN_LIFECYCLE,
    CAPABILITY_DATAFLOW_RUN_WORKER,
    DataflowRunDescriptor,
    DataflowRunLifecycleProvider,
    DataflowRunWorker,
    dataflow_run_lifecycle,
    dataflow_run_worker,
)
from govoplan_core.core.modules import ModuleContext, ModuleManifest
from govoplan_core.core.registry import PlatformRegistry


class _Provider:
    def start_run(self, session, principal, *, request):
        del session, principal
        return DataflowRunDescriptor(
            ref="dataflow-run:1",
            pipeline_ref=request.pipeline_ref,
            revision=request.revision,
            status="succeeded",
            definition_hash="abc",
            executor_version="test",
        )

    def get_run(self, session, principal, *, run_ref):
        del session, principal
        if run_ref != "dataflow-run:1":
            return None
        return DataflowRunDescriptor(
            ref=run_ref,
            pipeline_ref="pipeline:1",
            revision=1,
            status="succeeded",
            definition_hash="abc",
            executor_version="test",
        )

    def cancel_run(self, session, principal, *, run_ref):
        del session, principal
        return DataflowRunDescriptor(
            ref=run_ref,
            pipeline_ref="pipeline:1",
            revision=1,
            status="cancelled",
            definition_hash="abc",
            executor_version="test",
        )


class _Worker:
    def dispatch_pending(
        self,
        session,
        *,
        now=None,
        limit=10,
        worker_id=None,
    ):
        return {"claimed": 0}

    def purge_expired(self, session, *, now=None, limit=500):
        return {"purged": 0}


class DataflowContractTests(unittest.TestCase):
    def test_run_lifecycle_is_runtime_checkable_and_resolved(self) -> None:
        provider = _Provider()
        worker = _Worker()
        self.assertIsInstance(provider, DataflowRunLifecycleProvider)
        self.assertIsInstance(worker, DataflowRunWorker)
        registry = PlatformRegistry()
        registry.register(
            ModuleManifest(
                id="dataflow_contract_test",
                name="Dataflow contract test",
                version="test",
                capability_factories={
                    CAPABILITY_DATAFLOW_RUN_LIFECYCLE: lambda context: provider,
                    CAPABILITY_DATAFLOW_RUN_WORKER: lambda context: worker,
                },
            )
        )
        registry.configure_capability_context(
            ModuleContext(registry=registry, settings=object())
        )

        self.assertIs(provider, dataflow_run_lifecycle(registry))
        self.assertIs(worker, dataflow_run_worker(registry))
        self.assertIsNone(dataflow_run_lifecycle(PlatformRegistry()))
        self.assertIsNone(dataflow_run_worker(PlatformRegistry()))


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