from __future__ import annotations

from abstractruntime.core.models import EffectType, RunState, RunStatus
from abstractruntime.visualflow_compiler import compile_visualflow


def test_visualflow_compiler_supports_memory_compact_node() -> None:
    flow = {
        "id": "vf_memory_compact",
        "name": "vf_memory_compact",
        "entryNode": "start",
        "nodes": [
            {"id": "start", "type": "on_flow_start", "data": {"nodeType": "on_flow_start"}},
            {"id": "compact", "type": "memory_compact", "data": {"nodeType": "memory_compact"}},
        ],
        "edges": [
            {"source": "start", "sourceHandle": "exec-out", "target": "compact", "targetHandle": "exec-in"},
        ],
    }

    spec = compile_visualflow(flow)
    assert "compact" in spec.nodes

    run = RunState(run_id="run", workflow_id=str(spec.workflow_id), status=RunStatus.RUNNING, current_node="compact", vars={"_temp": {}})
    plan = spec.nodes["compact"](run, None)
    assert plan.effect is not None
    assert plan.effect.type == EffectType.MEMORY_COMPACT


def test_visualflow_compiler_supports_memory_tag_node() -> None:
    flow = {
        "id": "vf_memory_tag",
        "name": "vf_memory_tag",
        "entryNode": "start",
        "nodes": [
            {"id": "start", "type": "on_flow_start", "data": {"nodeType": "on_flow_start"}},
            {"id": "tag", "type": "memory_tag", "data": {"nodeType": "memory_tag"}},
        ],
        "edges": [
            {"source": "start", "sourceHandle": "exec-out", "target": "tag", "targetHandle": "exec-in"},
        ],
    }

    spec = compile_visualflow(flow)
    assert "tag" in spec.nodes

    run = RunState(run_id="run", workflow_id=str(spec.workflow_id), status=RunStatus.RUNNING, current_node="tag", vars={"_temp": {}})
    plan = spec.nodes["tag"](run, None)
    assert plan.effect is not None
    assert plan.effect.type == EffectType.MEMORY_TAG


def test_visualflow_compiler_supports_memory_kg_query_node() -> None:
    flow = {
        "id": "vf_memory_kg_query",
        "name": "vf_memory_kg_query",
        "entryNode": "start",
        "nodes": [
            {"id": "start", "type": "on_flow_start", "data": {"nodeType": "on_flow_start"}},
            {"id": "kgq", "type": "memory_kg_query", "data": {"nodeType": "memory_kg_query"}},
        ],
        "edges": [
            {"source": "start", "sourceHandle": "exec-out", "target": "kgq", "targetHandle": "exec-in"},
        ],
    }

    spec = compile_visualflow(flow)
    assert "kgq" in spec.nodes

    run = RunState(run_id="run", workflow_id=str(spec.workflow_id), status=RunStatus.RUNNING, current_node="kgq", vars={"_temp": {}})
    plan = spec.nodes["kgq"](run, None)
    assert plan.effect is not None
    assert plan.effect.type == EffectType.MEMORY_KG_QUERY


def test_visualflow_compiler_supports_memory_kg_resolve_node() -> None:
    flow = {
        "id": "vf_memory_kg_resolve",
        "name": "vf_memory_kg_resolve",
        "entryNode": "start",
        "nodes": [
            {"id": "start", "type": "on_flow_start", "data": {"nodeType": "on_flow_start"}},
            {"id": "kgr", "type": "memory_kg_resolve", "data": {"nodeType": "memory_kg_resolve"}},
        ],
        "edges": [
            {"source": "start", "sourceHandle": "exec-out", "target": "kgr", "targetHandle": "exec-in"},
        ],
    }

    spec = compile_visualflow(flow)
    assert "kgr" in spec.nodes

    run = RunState(run_id="run", workflow_id=str(spec.workflow_id), status=RunStatus.RUNNING, current_node="kgr", vars={"_temp": {}})
    plan = spec.nodes["kgr"](run, None)
    assert plan.effect is not None
    assert plan.effect.type == EffectType.MEMORY_KG_RESOLVE


def test_visualflow_memory_kg_resolve_maps_label_and_expected_type() -> None:
    flow = {
        "id": "vf_memory_kg_resolve_maps",
        "name": "vf_memory_kg_resolve_maps",
        "entryNode": "start",
        "nodes": [
            {"id": "start", "type": "on_flow_start", "data": {"nodeType": "on_flow_start"}},
            {"id": "kgr", "type": "memory_kg_resolve", "data": {"nodeType": "memory_kg_resolve"}},
        ],
        "edges": [
            {"source": "start", "sourceHandle": "exec-out", "target": "kgr", "targetHandle": "exec-in"},
        ],
    }

    spec = compile_visualflow(flow)
    run = RunState(
        run_id="run",
        workflow_id=str(spec.workflow_id),
        status=RunStatus.RUNNING,
        current_node="kgr",
        vars={"_temp": {}, "_last_output": {"label": "doctor noonien soong", "expected_type": "schema:person", "scope": "global"}},
    )
    plan = spec.nodes["kgr"](run, None)
    assert plan.effect is not None
    assert plan.effect.type == EffectType.MEMORY_KG_RESOLVE
    assert isinstance(plan.effect.payload, dict)
    assert plan.effect.payload.get("label") == "doctor noonien soong"
    assert plan.effect.payload.get("expected_type") == "schema:person"
    assert plan.effect.payload.get("scope") == "global"


def test_visualflow_memory_kg_query_maps_min_score() -> None:
    flow = {
        "id": "vf_memory_kg_query_min_score",
        "name": "vf_memory_kg_query_min_score",
        "entryNode": "start",
        "nodes": [
            {"id": "start", "type": "on_flow_start", "data": {"nodeType": "on_flow_start"}},
            {"id": "kgq", "type": "memory_kg_query", "data": {"nodeType": "memory_kg_query"}},
        ],
        "edges": [
            {"source": "start", "sourceHandle": "exec-out", "target": "kgq", "targetHandle": "exec-in"},
        ],
    }

    spec = compile_visualflow(flow)
    run = RunState(
        run_id="run",
        workflow_id=str(spec.workflow_id),
        status=RunStatus.RUNNING,
        current_node="kgq",
        # For VisualFlow-compiled effect nodes, resolved pin values are passed
        # through the visual executor via `_last_output`.
        vars={"_temp": {}, "_last_output": {"query_text": "android", "min_score": 0.33}},
    )
    plan = spec.nodes["kgq"](run, None)
    assert plan.effect is not None
    assert plan.effect.type == EffectType.MEMORY_KG_QUERY
    assert isinstance(plan.effect.payload, dict)
    assert plan.effect.payload.get("query_text") == "android"
    assert plan.effect.payload.get("min_score") == 0.33


def test_visualflow_compiler_supports_memory_kg_assert_node() -> None:
    flow = {
        "id": "vf_memory_kg_assert",
        "name": "vf_memory_kg_assert",
        "entryNode": "start",
        "nodes": [
            {"id": "start", "type": "on_flow_start", "data": {"nodeType": "on_flow_start"}},
            {"id": "kga", "type": "memory_kg_assert", "data": {"nodeType": "memory_kg_assert"}},
        ],
        "edges": [
            {"source": "start", "sourceHandle": "exec-out", "target": "kga", "targetHandle": "exec-in"},
        ],
    }

    spec = compile_visualflow(flow)
    assert "kga" in spec.nodes

    run = RunState(run_id="run", workflow_id=str(spec.workflow_id), status=RunStatus.RUNNING, current_node="kga", vars={"_temp": {}})
    plan = spec.nodes["kga"](run, None)
    assert plan.effect is not None
    assert plan.effect.type == EffectType.MEMORY_KG_ASSERT


def test_visualflow_memory_note_maps_keep_in_context_and_location() -> None:
    flow = {
        "id": "vf_memory_note_mapping",
        "name": "vf_memory_note_mapping",
        "entryNode": "start",
        "nodes": [
            {
                "id": "start",
                "type": "on_flow_start",
                "data": {
                    "nodeType": "on_flow_start",
                    "outputs": [
                        {"id": "exec-out", "type": "execution"},
                        {"id": "content", "type": "string"},
                        {"id": "keep_in_context", "type": "boolean"},
                        {"id": "location", "type": "string"},
                    ],
                    "pinDefaults": {"content": "hello", "keep_in_context": True, "location": "flow:test"},
                },
            },
            {"id": "note", "type": "memory_note", "data": {"nodeType": "memory_note"}},
        ],
        "edges": [
            {"source": "start", "sourceHandle": "exec-out", "target": "note", "targetHandle": "exec-in"},
            {"source": "start", "sourceHandle": "content", "target": "note", "targetHandle": "content"},
            {"source": "start", "sourceHandle": "keep_in_context", "target": "note", "targetHandle": "keep_in_context"},
            {"source": "start", "sourceHandle": "location", "target": "note", "targetHandle": "location"},
        ],
    }

    spec = compile_visualflow(flow)
    assert "note" in spec.nodes

    run = RunState(
        run_id="run",
        workflow_id=str(spec.workflow_id),
        status=RunStatus.RUNNING,
        current_node="start",
        vars={"_temp": {}},
    )
    plan = spec.nodes["start"](run, None)
    assert plan.next_node == "note"
    run.current_node = "note"

    plan = spec.nodes["note"](run, None)
    assert plan.effect is not None
    assert plan.effect.type == EffectType.MEMORY_NOTE
    assert plan.effect.payload.get("keep_in_context") is True
    assert plan.effect.payload.get("location") == "flow:test"


def test_visualflow_memory_query_maps_tags_mode_usernames_and_locations() -> None:
    flow = {
        "id": "vf_memory_query_mapping",
        "name": "vf_memory_query_mapping",
        "entryNode": "start",
        "nodes": [
            {
                "id": "start",
                "type": "on_flow_start",
                "data": {
                    "nodeType": "on_flow_start",
                    "outputs": [
                        {"id": "exec-out", "type": "execution"},
                        {"id": "query", "type": "string"},
                        {"id": "limit", "type": "number"},
                        {"id": "recall_level", "type": "string"},
                        {"id": "tags_mode", "type": "string"},
                        {"id": "usernames", "type": "array"},
                        {"id": "locations", "type": "array"},
                    ],
                    "pinDefaults": {
                        "query": "who",
                        "limit": 3,
                        "recall_level": "urgent",
                        "tags_mode": "any",
                        "usernames": ["alice"],
                        "locations": ["flow:x"],
                    },
                },
            },
            {"id": "query", "type": "memory_query", "data": {"nodeType": "memory_query"}},
        ],
        "edges": [
            {"source": "start", "sourceHandle": "exec-out", "target": "query", "targetHandle": "exec-in"},
            {"source": "start", "sourceHandle": "query", "target": "query", "targetHandle": "query"},
            {"source": "start", "sourceHandle": "limit", "target": "query", "targetHandle": "limit"},
            {"source": "start", "sourceHandle": "recall_level", "target": "query", "targetHandle": "recall_level"},
            {"source": "start", "sourceHandle": "tags_mode", "target": "query", "targetHandle": "tags_mode"},
            {"source": "start", "sourceHandle": "usernames", "target": "query", "targetHandle": "usernames"},
            {"source": "start", "sourceHandle": "locations", "target": "query", "targetHandle": "locations"},
        ],
    }

    spec = compile_visualflow(flow)
    assert "query" in spec.nodes

    run = RunState(
        run_id="run",
        workflow_id=str(spec.workflow_id),
        status=RunStatus.RUNNING,
        current_node="start",
        vars={"_temp": {}},
    )
    plan = spec.nodes["start"](run, None)
    assert plan.next_node == "query"
    run.current_node = "query"

    plan = spec.nodes["query"](run, None)
    assert plan.effect is not None
    assert plan.effect.type == EffectType.MEMORY_QUERY
    assert plan.effect.payload.get("recall_level") == "urgent"
    assert plan.effect.payload.get("tags_mode") == "any"
    assert plan.effect.payload.get("usernames") == ["alice"]
    assert plan.effect.payload.get("locations") == ["flow:x"]


def test_visualflow_memory_kg_query_maps_recall_level() -> None:
    flow = {
        "id": "vf_memory_kg_query_recall_level",
        "name": "vf_memory_kg_query_recall_level",
        "entryNode": "start",
        "nodes": [
            {"id": "start", "type": "on_flow_start", "data": {"nodeType": "on_flow_start"}},
            {"id": "kgq", "type": "memory_kg_query", "data": {"nodeType": "memory_kg_query"}},
        ],
        "edges": [
            {"source": "start", "sourceHandle": "exec-out", "target": "kgq", "targetHandle": "exec-in"},
        ],
    }

    spec = compile_visualflow(flow)
    run = RunState(
        run_id="run",
        workflow_id=str(spec.workflow_id),
        status=RunStatus.RUNNING,
        current_node="kgq",
        vars={"_temp": {}, "_last_output": {"query_text": "android", "min_score": 0.33, "recall_level": "standard"}},
    )
    plan = spec.nodes["kgq"](run, None)
    assert plan.effect is not None
    assert plan.effect.type == EffectType.MEMORY_KG_QUERY
    assert plan.effect.payload.get("query_text") == "android"
    assert plan.effect.payload.get("min_score") == 0.33
    assert plan.effect.payload.get("recall_level") == "standard"


def test_visualflow_memory_rehydrate_maps_recall_level() -> None:
    flow = {
        "id": "vf_memory_rehydrate_recall_level",
        "name": "vf_memory_rehydrate_recall_level",
        "entryNode": "start",
        "nodes": [
            {
                "id": "start",
                "type": "on_flow_start",
                "data": {
                    "nodeType": "on_flow_start",
                    "outputs": [
                        {"id": "exec-out", "type": "execution"},
                        {"id": "span_ids", "type": "array"},
                        {"id": "max_messages", "type": "number"},
                        {"id": "recall_level", "type": "string"},
                    ],
                    "pinDefaults": {
                        "span_ids": ["a1", "a2"],
                        "max_messages": 25,
                        "recall_level": "urgent",
                    },
                },
            },
            {"id": "rehydrate", "type": "memory_rehydrate", "data": {"nodeType": "memory_rehydrate"}},
        ],
        "edges": [
            {"source": "start", "sourceHandle": "exec-out", "target": "rehydrate", "targetHandle": "exec-in"},
            {"source": "start", "sourceHandle": "span_ids", "target": "rehydrate", "targetHandle": "span_ids"},
            {"source": "start", "sourceHandle": "max_messages", "target": "rehydrate", "targetHandle": "max_messages"},
            {"source": "start", "sourceHandle": "recall_level", "target": "rehydrate", "targetHandle": "recall_level"},
        ],
    }

    spec = compile_visualflow(flow)
    assert "rehydrate" in spec.nodes

    run = RunState(
        run_id="run",
        workflow_id=str(spec.workflow_id),
        status=RunStatus.RUNNING,
        current_node="start",
        vars={"_temp": {}},
    )
    plan = spec.nodes["start"](run, None)
    assert plan.next_node == "rehydrate"
    run.current_node = "rehydrate"

    plan2 = spec.nodes["rehydrate"](run, None)
    assert plan2.effect is not None
    assert plan2.effect.type == EffectType.MEMORY_REHYDRATE
    assert plan2.effect.payload.get("span_ids") == ["a1", "a2"]
    assert plan2.effect.payload.get("max_messages") == 25
    assert plan2.effect.payload.get("recall_level") == "urgent"
