"""Tests for metaflow_extensions.prefect.plugins.prefect._graph."""
from __future__ import annotations

from typing import Any
from unittest.mock import MagicMock

import pytest

from metaflow_extensions.prefect.plugins.prefect._graph import analyze_graph
from metaflow_extensions.prefect.plugins.prefect._types import (
    FlowSpec,
    NodeType,
    TriggerOnFinishSpec,
    TriggerSpec,
)
from metaflow_extensions.prefect.plugins.prefect.exception import NotSupportedException


class TestAnalyzeGraphSimple:
    """analyze_graph with the simple linear flow."""

    def test_returns_flow_spec(self, simple_flow_graph: tuple[Any, Any]) -> None:
        """analyze_graph returns a FlowSpec."""
        graph, flow = simple_flow_graph
        spec = analyze_graph(graph, flow)
        assert isinstance(spec, FlowSpec)

    def test_flow_name(self, simple_flow_graph: tuple[Any, Any]) -> None:
        """FlowSpec.name matches the flow class name."""
        graph, flow = simple_flow_graph
        spec = analyze_graph(graph, flow)
        assert spec.name == "SimpleFlow"

    def test_topological_order_starts_with_start(self, simple_flow_graph: tuple[Any, Any]) -> None:
        """First step in topological order is always 'start'."""
        graph, flow = simple_flow_graph
        spec = analyze_graph(graph, flow)
        assert spec.steps[0].name == "start"

    def test_topological_order_ends_with_end(self, simple_flow_graph: tuple[Any, Any]) -> None:
        """Last step in topological order is always 'end'."""
        graph, flow = simple_flow_graph
        spec = analyze_graph(graph, flow)
        assert spec.steps[-1].name == "end"

    def test_all_steps_present(self, simple_flow_graph: tuple[Any, Any]) -> None:
        """All three steps (start, process, end) appear in the spec."""
        graph, flow = simple_flow_graph
        spec = analyze_graph(graph, flow)
        names = {s.name for s in spec.steps}
        assert names == {"start", "process", "end"}

    def test_start_node_type(self, simple_flow_graph: tuple[Any, Any]) -> None:
        """The start step has NodeType.START."""
        graph, flow = simple_flow_graph
        spec = analyze_graph(graph, flow)
        start = next(s for s in spec.steps if s.name == "start")
        assert start.node_type == NodeType.START

    def test_end_node_type(self, simple_flow_graph: tuple[Any, Any]) -> None:
        """The end step has NodeType.END."""
        graph, flow = simple_flow_graph
        spec = analyze_graph(graph, flow)
        end = next(s for s in spec.steps if s.name == "end")
        assert end.node_type == NodeType.END

    def test_no_parameters(self, simple_flow_graph: tuple[Any, Any]) -> None:
        """SimpleFlow has no declared parameters."""
        graph, flow = simple_flow_graph
        spec = analyze_graph(graph, flow)
        assert spec.parameters == ()


class TestAnalyzeGraphBranch:
    """analyze_graph with the branch/join flow."""

    def test_split_node_type(self, branch_flow_graph: tuple[Any, Any]) -> None:
        """The start step has type SPLIT when it fans out."""
        graph, flow = branch_flow_graph
        spec = analyze_graph(graph, flow)
        start = next(s for s in spec.steps if s.name == "start")
        # In Metaflow, a step that calls self.next(a, b) has type "split"
        assert start.node_type == NodeType.SPLIT

    def test_join_flags(self, branch_flow_graph: tuple[Any, Any]) -> None:
        """The join step is detected as a split-join (not a foreach-join)."""
        graph, flow = branch_flow_graph
        spec = analyze_graph(graph, flow)
        join = next(s for s in spec.steps if s.name == "join")
        assert join.is_split_join is True
        assert join.is_foreach_join is False

    def test_join_in_funcs(self, branch_flow_graph: tuple[Any, Any]) -> None:
        """The join step lists both branches as in_funcs."""
        graph, flow = branch_flow_graph
        spec = analyze_graph(graph, flow)
        join = next(s for s in spec.steps if s.name == "join")
        assert set(join.in_funcs) == {"branch_a", "branch_b"}

    def test_all_steps_present(self, branch_flow_graph: tuple[Any, Any]) -> None:
        graph, flow = branch_flow_graph
        spec = analyze_graph(graph, flow)
        names = {s.name for s in spec.steps}
        assert names == {"start", "branch_a", "branch_b", "join", "end"}


class TestAnalyzeGraphForeach:
    """analyze_graph with the foreach flow."""

    def test_foreach_step_type(self, foreach_flow_graph: tuple[Any, Any]) -> None:
        """The start step has type FOREACH when it fans out with foreach=."""
        graph, flow = foreach_flow_graph
        spec = analyze_graph(graph, flow)
        start = next(s for s in spec.steps if s.name == "start")
        assert start.node_type == NodeType.FOREACH

    def test_foreach_join_flag(self, foreach_flow_graph: tuple[Any, Any]) -> None:
        """The join step that closes a foreach is detected correctly."""
        graph, flow = foreach_flow_graph
        spec = analyze_graph(graph, flow)
        join = next(s for s in spec.steps if s.name == "join_step")
        assert join.is_foreach_join is True
        assert join.is_split_join is False


class TestAnalyzeGraphParams:
    """analyze_graph with the parametrised flow."""

    def test_parameters_extracted(self, param_flow_graph: tuple[Any, Any]) -> None:
        """Parameters are present in the FlowSpec."""
        graph, flow = param_flow_graph
        spec = analyze_graph(graph, flow)
        param_names = {p.name for p in spec.parameters}
        assert "message" in param_names
        assert "count" in param_names

    def test_parameter_defaults(self, param_flow_graph: tuple[Any, Any]) -> None:
        """Default values are evaluated at analysis time."""
        graph, flow = param_flow_graph
        spec = analyze_graph(graph, flow)
        msg_param = next(p for p in spec.parameters if p.name == "message")
        count_param = next(p for p in spec.parameters if p.name == "count")
        assert msg_param.default == "hello"
        assert count_param.default == 3


def _make_mock_graph(parallel_foreach: bool = False, deco_names: list[str] | None = None) -> Any:
    """Return a (graph, flow) mock pair for validation tests."""
    from unittest.mock import MagicMock

    deco_names = deco_names or []
    decos = []
    for dname in deco_names:
        d = MagicMock()
        d.name = dname
        decos.append(d)

    node = MagicMock()
    node.name = "start"
    node.parallel_foreach = parallel_foreach
    node.decorators = decos
    node.in_funcs = []
    node.out_funcs = []
    node.split_parents = []
    node.type = "start"

    graph = MagicMock()
    graph.__iter__ = MagicMock(return_value=iter([node]))
    graph.__getitem__ = MagicMock(return_value=node)

    flow = MagicMock()
    flow.name = "BadFlow"
    flow._get_parameters = MagicMock(return_value=[])
    flow._flow_decorators = {}

    return graph, flow


def _make_2level_nested_foreach_graph() -> tuple[Any, Any]:
    """2-level nesting: foreach step inside another foreach (1 foreach ancestor).

    This should NOT raise — 2-level nesting is supported.
    """
    outer = MagicMock()
    outer.name = "start"
    outer.type = "foreach"
    outer.parallel_foreach = False
    outer.decorators = []
    outer.in_funcs = []
    outer.out_funcs = ["inner_foreach"]
    outer.split_parents = []

    inner = MagicMock()
    inner.name = "inner_foreach"
    inner.type = "foreach"
    inner.parallel_foreach = False
    inner.decorators = []
    inner.in_funcs = ["start"]
    inner.out_funcs = []
    inner.split_parents = ["start"]  # 1 foreach ancestor — allowed

    graph = MagicMock()
    graph.__iter__ = MagicMock(return_value=iter([outer, inner]))
    graph.__getitem__ = MagicMock(
        side_effect=lambda name: outer if name == "start" else inner
    )

    flow = MagicMock()
    flow.name = "TwoLevelForeach"
    flow._get_parameters = MagicMock(return_value=[])
    flow._flow_decorators = {}
    return graph, flow


def _make_3level_nested_foreach_graph() -> tuple[Any, Any]:
    """3-level nesting: foreach inside foreach inside foreach (2 foreach ancestors)."""
    outer = MagicMock()
    outer.name = "start"
    outer.type = "foreach"
    outer.parallel_foreach = False
    outer.decorators = []
    outer.in_funcs = []
    outer.out_funcs = ["middle"]
    outer.split_parents = []

    middle = MagicMock()
    middle.name = "middle"
    middle.type = "foreach"
    middle.parallel_foreach = False
    middle.decorators = []
    middle.in_funcs = ["start"]
    middle.out_funcs = ["deepest"]
    middle.split_parents = ["start"]

    deepest = MagicMock()
    deepest.name = "deepest"
    deepest.type = "foreach"
    deepest.parallel_foreach = False
    deepest.decorators = []
    deepest.in_funcs = ["middle"]
    deepest.out_funcs = []
    deepest.split_parents = ["start", "middle"]  # 2 foreach ancestors — raises

    def _getitem(name: str) -> MagicMock:
        if name == "start":
            return outer
        if name == "middle":
            return middle
        return deepest

    graph = MagicMock()
    graph.__iter__ = MagicMock(return_value=iter([outer, middle, deepest]))
    graph.__getitem__ = MagicMock(side_effect=_getitem)

    flow = MagicMock()
    flow.name = "ThreeLevelForeach"
    flow._get_parameters = MagicMock(return_value=[])
    flow._flow_decorators = {}
    return graph, flow


class TestValidation:
    """analyze_graph raises for unsupported features."""

    def test_parallel_foreach_raises(self) -> None:
        graph, flow = _make_mock_graph(parallel_foreach=True)
        with pytest.raises(NotSupportedException):
            analyze_graph(graph, flow)

    def test_batch_raises(self) -> None:
        graph, flow = _make_mock_graph(deco_names=["batch"])
        with pytest.raises(NotSupportedException, match="@batch"):
            analyze_graph(graph, flow)

    def test_slurm_raises(self) -> None:
        graph, flow = _make_mock_graph(deco_names=["slurm"])
        with pytest.raises(NotSupportedException, match="@slurm"):
            analyze_graph(graph, flow)

    def test_trigger_does_not_raise(self) -> None:
        graph, flow = _make_mock_graph()
        trigger = MagicMock()
        trigger.triggers = [{"name": "my_event", "parameters": {}}]
        flow._flow_decorators = {"trigger": [trigger]}
        # Should not raise — @trigger is handled via Prefect automations.
        analyze_graph(graph, flow)

    def test_trigger_extracted_into_spec(self) -> None:
        graph, flow = _make_mock_graph()
        trigger = MagicMock()
        trigger.triggers = [{"name": "data.ready", "parameters": {"x": "y"}}]
        flow._flow_decorators = {"trigger": [trigger]}
        spec = analyze_graph(graph, flow)
        assert len(spec.triggers) == 1
        assert spec.triggers[0] == TriggerSpec(
            event_name="data.ready", parameter_map=(("x", "y"),)
        )

    def test_trigger_on_finish_does_not_raise(self) -> None:
        graph, flow = _make_mock_graph()
        t = MagicMock()
        t.triggers = [{"fq_name": "OtherFlow", "flow": "OtherFlow"}]
        flow._flow_decorators = {"trigger_on_finish": [t]}
        # Should not raise — @trigger_on_finish is handled via Prefect automations.
        analyze_graph(graph, flow)

    def test_trigger_on_finish_extracted_into_spec(self) -> None:
        graph, flow = _make_mock_graph()
        t = MagicMock()
        t.triggers = [{"fq_name": "OtherFlow", "flow": "OtherFlow"}]
        flow._flow_decorators = {"trigger_on_finish": [t]}
        spec = analyze_graph(graph, flow)
        assert len(spec.trigger_on_finishes) == 1
        assert spec.trigger_on_finishes[0] == TriggerOnFinishSpec(flow_name="OtherFlow")

    def test_exit_hook_raises(self) -> None:
        graph, flow = _make_mock_graph()
        e = MagicMock()
        flow._flow_decorators = {"exit_hook": [e]}
        with pytest.raises(NotSupportedException, match="@exit_hook"):
            analyze_graph(graph, flow)

    def test_nested_foreach_does_not_raise(self) -> None:
        """2-level nested foreach is allowed (arbitrary nesting depth supported)."""
        graph, flow = _make_2level_nested_foreach_graph()
        try:
            analyze_graph(graph, flow)
        except NotSupportedException:
            pytest.fail("2-level nested foreach should not raise NotSupportedException")
        except Exception:
            pass  # Other errors from incomplete mock graph are OK

    def test_3level_nested_foreach_does_not_raise(self) -> None:
        """3-level nested foreach is also allowed — no depth limit."""
        graph, flow = _make_3level_nested_foreach_graph()
        try:
            analyze_graph(graph, flow)
        except NotSupportedException:
            pytest.fail("3-level nested foreach should not raise NotSupportedException")
        except Exception:
            pass  # Other errors from incomplete mock graph are OK


class TestResourcesExtraction:
    """StepSpec populated from @resources decorator."""

    def test_cpu_extracted(self, resources_flow_graph: tuple[Any, Any]) -> None:
        graph, flow = resources_flow_graph
        spec = analyze_graph(graph, flow)
        start = next(s for s in spec.steps if s.name == "start")
        assert start.resource_cpu == 4

    def test_memory_extracted(self, resources_flow_graph: tuple[Any, Any]) -> None:
        graph, flow = resources_flow_graph
        spec = analyze_graph(graph, flow)
        start = next(s for s in spec.steps if s.name == "start")
        assert start.resource_memory == 8192

    def test_gpu_extracted(self, resources_flow_graph: tuple[Any, Any]) -> None:
        graph, flow = resources_flow_graph
        spec = analyze_graph(graph, flow)
        end = next(s for s in spec.steps if s.name == "end")
        assert end.resource_gpu == 1

    def test_no_resources_on_plain_step(self, simple_flow_graph: tuple[Any, Any]) -> None:
        graph, flow = simple_flow_graph
        spec = analyze_graph(graph, flow)
        for step in spec.steps:
            assert step.resource_cpu is None
            assert step.resource_gpu is None
            assert step.resource_memory is None


class TestRequiredParams:
    """Required parameters (no default) handled correctly."""

    def test_required_param_has_required_flag(
        self, required_param_flow_graph: tuple[Any, Any]
    ) -> None:
        graph, flow = required_param_flow_graph
        spec = analyze_graph(graph, flow)
        msg = next(p for p in spec.parameters if p.name == "message")
        assert msg.required is True

    def test_required_param_default_is_none(
        self, required_param_flow_graph: tuple[Any, Any]
    ) -> None:
        graph, flow = required_param_flow_graph
        spec = analyze_graph(graph, flow)
        msg = next(p for p in spec.parameters if p.name == "message")
        assert msg.default is None

    def test_optional_param_not_required(
        self, required_param_flow_graph: tuple[Any, Any]
    ) -> None:
        graph, flow = required_param_flow_graph
        spec = analyze_graph(graph, flow)
        count = next(p for p in spec.parameters if p.name == "count")
        assert count.required is False
        assert count.default == 5

    def test_required_param_type_inferred(
        self, required_param_flow_graph: tuple[Any, Any]
    ) -> None:
        graph, flow = required_param_flow_graph
        spec = analyze_graph(graph, flow)
        msg = next(p for p in spec.parameters if p.name == "message")
        assert msg.type_name == "str"


class TestDecoratorExtraction:
    """StepSpec fields populated from @retry, @timeout, @environment."""

    def test_timeout_seconds_from_timeout_deco(self, decorator_flow_graph: tuple[Any, Any]) -> None:
        graph, flow = decorator_flow_graph
        spec = analyze_graph(graph, flow)
        start = next(s for s in spec.steps if s.name == "start")
        assert start.timeout_seconds == 300

    def test_timeout_minutes_from_timeout_deco(self, decorator_flow_graph: tuple[Any, Any]) -> None:
        graph, flow = decorator_flow_graph
        spec = analyze_graph(graph, flow)
        end = next(s for s in spec.steps if s.name == "end")
        assert end.timeout_seconds == 300  # 5 * 60

    def test_retry_delay_seconds(self, decorator_flow_graph: tuple[Any, Any]) -> None:
        graph, flow = decorator_flow_graph
        spec = analyze_graph(graph, flow)
        start = next(s for s in spec.steps if s.name == "start")
        assert start.retry_delay_seconds == 60  # 1 minute

    def test_retry_count(self, decorator_flow_graph: tuple[Any, Any]) -> None:
        graph, flow = decorator_flow_graph
        spec = analyze_graph(graph, flow)
        start = next(s for s in spec.steps if s.name == "start")
        assert start.max_user_code_retries == 2

    def test_env_vars_extracted(self, decorator_flow_graph: tuple[Any, Any]) -> None:
        graph, flow = decorator_flow_graph
        spec = analyze_graph(graph, flow)
        start = next(s for s in spec.steps if s.name == "start")
        env = dict(start.env_vars)
        assert env.get("MY_VAR") == "hello"
        assert env.get("OTHER") == "world"

    def test_no_env_vars_on_plain_step(self, decorator_flow_graph: tuple[Any, Any]) -> None:
        graph, flow = decorator_flow_graph
        spec = analyze_graph(graph, flow)
        end = next(s for s in spec.steps if s.name == "end")
        assert end.env_vars == ()

    def test_no_timeout_on_plain_step(self, simple_flow_graph: tuple[Any, Any]) -> None:
        graph, flow = simple_flow_graph
        spec = analyze_graph(graph, flow)
        for step in spec.steps:
            assert step.timeout_seconds is None

    def test_no_retry_delay_on_plain_step(self, simple_flow_graph: tuple[Any, Any]) -> None:
        graph, flow = simple_flow_graph
        spec = analyze_graph(graph, flow)
        for step in spec.steps:
            assert step.retry_delay_seconds is None


class TestValidationAdversarial:
    """Edge-case and adversarial validation tests targeting uncovered branches."""

    def test_condition_raises(self) -> None:
        """@condition decorator raises NotSupportedException."""
        graph, flow = _make_mock_graph(deco_names=["condition"])
        with pytest.raises(NotSupportedException, match="@condition"):
            analyze_graph(graph, flow)

    def test_unknown_node_type_falls_back_to_linear(self) -> None:
        """An unrecognised node type string falls back to NodeType.LINEAR without error."""
        from metaflow_extensions.prefect.plugins.prefect._graph import _topological_order

        node = MagicMock()
        node.name = "start"
        node.type = "future_metaflow_node_type"  # not in NodeType enum
        node.in_funcs = []
        node.out_funcs = []
        node.split_parents = []
        node.decorators = []

        graph = MagicMock()
        graph.__iter__ = MagicMock(return_value=iter([node]))
        graph.__getitem__ = MagicMock(return_value=node)

        steps = _topological_order(graph)
        assert len(steps) == 1
        assert steps[0].node_type == NodeType.LINEAR

    def test_schedule_cron_passthrough(self) -> None:
        """@schedule(cron=...) passes the cron string through unchanged."""
        from metaflow_extensions.prefect.plugins.prefect._graph import _extract_schedule

        s = MagicMock()
        s.attributes = {"cron": "*/5 * * * *"}
        flow = MagicMock()
        flow._flow_decorators = MagicMock()
        flow._flow_decorators.get = MagicMock(return_value=[s])
        assert _extract_schedule(flow) == "*/5 * * * *"

    def test_schedule_weekly(self) -> None:
        """@schedule(weekly=True) produces weekly cron string."""
        from metaflow_extensions.prefect.plugins.prefect._graph import _extract_schedule

        s = MagicMock()
        s.attributes = {"weekly": True}
        flow = MagicMock()
        flow._flow_decorators = MagicMock()
        flow._flow_decorators.get = MagicMock(return_value=[s])
        assert _extract_schedule(flow) == "0 0 * * 0"

    def test_schedule_hourly(self) -> None:
        """@schedule(hourly=True) produces hourly cron string."""
        from metaflow_extensions.prefect.plugins.prefect._graph import _extract_schedule

        s = MagicMock()
        s.attributes = {"hourly": True}
        flow = MagicMock()
        flow._flow_decorators = MagicMock()
        flow._flow_decorators.get = MagicMock(return_value=[s])
        assert _extract_schedule(flow) == "0 * * * *"

    def test_schedule_daily(self) -> None:
        """@schedule(daily=True) produces daily cron string."""
        from metaflow_extensions.prefect.plugins.prefect._graph import _extract_schedule

        s = MagicMock()
        s.attributes = {"daily": True}
        flow = MagicMock()
        flow._flow_decorators = MagicMock()
        flow._flow_decorators.get = MagicMock(return_value=[s])
        assert _extract_schedule(flow) == "0 0 * * *"

    def test_schedule_unknown_attrs_returns_none(self) -> None:
        """@schedule with unrecognised attrs returns None rather than crashing."""
        from metaflow_extensions.prefect.plugins.prefect._graph import _extract_schedule

        s = MagicMock()
        s.attributes = {"some_future_attr": True}
        flow = MagicMock()
        flow._flow_decorators = MagicMock()
        flow._flow_decorators.get = MagicMock(return_value=[s])
        assert _extract_schedule(flow) is None

    def test_project_name_extracted(self) -> None:
        """@project(name=...) populates FlowSpec.project_name."""
        graph, flow = _make_mock_graph()
        proj = MagicMock()
        proj.attributes = {"name": "my_project"}
        flow._flow_decorators = {"project": [proj]}
        spec = analyze_graph(graph, flow)
        assert spec.project_name == "my_project"

    def test_trigger_non_dict_entry_skipped(self) -> None:
        """A non-dict entry in the triggers list is silently skipped."""
        graph, flow = _make_mock_graph()
        trigger = MagicMock()
        trigger.triggers = ["not_a_dict", {"name": "ok_event"}]
        flow._flow_decorators = {"trigger": [trigger]}
        spec = analyze_graph(graph, flow)
        assert len(spec.triggers) == 1
        assert spec.triggers[0].event_name == "ok_event"

    def test_trigger_none_name_warns_and_skips(self) -> None:
        """A trigger entry with a None name emits a UserWarning and is skipped."""
        import warnings

        graph, flow = _make_mock_graph()
        trigger = MagicMock()
        trigger.triggers = [{"name": None}, {"name": "good_event"}]
        flow._flow_decorators = {"trigger": [trigger]}
        with warnings.catch_warnings(record=True) as caught:
            warnings.simplefilter("always")
            spec = analyze_graph(graph, flow)
        assert any("non-string" in str(w.message) for w in caught)
        # The None-name entry is skipped; only "good_event" survives.
        assert len(spec.triggers) == 1
        assert spec.triggers[0].event_name == "good_event"

    def test_trigger_on_finish_non_dict_entry_skipped(self) -> None:
        """A non-dict entry in trigger_on_finish triggers is silently skipped."""
        graph, flow = _make_mock_graph()
        t = MagicMock()
        t.triggers = [42, {"flow": "GoodFlow"}]
        flow._flow_decorators = {"trigger_on_finish": [t]}
        spec = analyze_graph(graph, flow)
        assert len(spec.trigger_on_finishes) == 1
        assert spec.trigger_on_finishes[0].flow_name == "GoodFlow"

    def test_trigger_on_finish_missing_name_warns_and_skips(self) -> None:
        """A trigger_on_finish entry with no flow name emits a UserWarning and is skipped."""
        import warnings

        graph, flow = _make_mock_graph()
        t = MagicMock()
        t.triggers = [{"flow": None}, {"flow": "RealFlow"}]
        flow._flow_decorators = {"trigger_on_finish": [t]}
        with warnings.catch_warnings(record=True) as caught:
            warnings.simplefilter("always")
            spec = analyze_graph(graph, flow)
        assert any("missing flow name" in str(w.message) for w in caught)
        assert len(spec.trigger_on_finishes) == 1
        assert spec.trigger_on_finishes[0].flow_name == "RealFlow"
