#
#  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

import asyncio
import importlib.util
import json
import sys
from pathlib import Path
from types import ModuleType, SimpleNamespace

import pytest


def _load_canvas_runtime(monkeypatch):
    repo_root = Path(__file__).resolve().parents[4]

    quart = ModuleType("quart")
    quart.make_response = lambda *a, **kw: None
    quart.jsonify = lambda *a, **kw: None
    monkeypatch.setitem(sys.modules, "quart", quart)

    common_pkg = ModuleType("common")
    common_pkg.__path__ = [str(repo_root / "common")]
    monkeypatch.setitem(sys.modules, "common", common_pkg)

    common_constants = ModuleType("common.constants")
    common_constants.LLMType = SimpleNamespace(TTS="tts")
    monkeypatch.setitem(sys.modules, "common.constants", common_constants)

    common_misc = ModuleType("common.misc_utils")
    common_misc.get_uuid = lambda: "uuid"
    common_misc.hash_str2int = lambda x: 1

    async def _thread_pool_exec(fn, *args, **kwargs):
        return fn(*args, **kwargs)

    common_misc.thread_pool_exec = _thread_pool_exec
    monkeypatch.setitem(sys.modules, "common.misc_utils", common_misc)

    common_conn = ModuleType("common.connection_utils")

    def timeout(_seconds):
        def decorator(fn):
            return fn

        return decorator

    common_conn.timeout = timeout
    monkeypatch.setitem(sys.modules, "common.connection_utils", common_conn)

    common_ex = ModuleType("common.exceptions")

    class TaskCanceledException(Exception):
        pass

    common_ex.TaskCanceledException = TaskCanceledException
    monkeypatch.setitem(sys.modules, "common.exceptions", common_ex)

    api_pkg = ModuleType("api")
    api_pkg.__path__ = [str(repo_root / "api")]
    monkeypatch.setitem(sys.modules, "api", api_pkg)
    api_db_pkg = ModuleType("api.db")
    api_db_pkg.__path__ = [str(repo_root / "api" / "db")]
    monkeypatch.setitem(sys.modules, "api.db", api_db_pkg)
    api_db_services_pkg = ModuleType("api.db.services")
    api_db_services_pkg.__path__ = [str(repo_root / "api" / "db" / "services")]
    monkeypatch.setitem(sys.modules, "api.db.services", api_db_services_pkg)
    api_db_joint_pkg = ModuleType("api.db.joint_services")
    api_db_joint_pkg.__path__ = [str(repo_root / "api" / "db" / "joint_services")]
    monkeypatch.setitem(sys.modules, "api.db.joint_services", api_db_joint_pkg)

    file_service = ModuleType("api.db.services.file_service")
    file_service.FileService = object
    monkeypatch.setitem(sys.modules, "api.db.services.file_service", file_service)

    llm_service = ModuleType("api.db.services.llm_service")
    llm_service.LLMBundle = object
    monkeypatch.setitem(sys.modules, "api.db.services.llm_service", llm_service)

    task_service = ModuleType("api.db.services.task_service")
    task_service.has_canceled = lambda _task_id: False
    monkeypatch.setitem(sys.modules, "api.db.services.task_service", task_service)

    tenant_model_service = ModuleType("api.db.joint_services.tenant_model_service")
    tenant_model_service.get_tenant_default_model_by_type = lambda *_a, **_kw: None
    monkeypatch.setitem(
        sys.modules,
        "api.db.joint_services.tenant_model_service",
        tenant_model_service,
    )

    rag_pkg = ModuleType("rag")
    rag_pkg.__path__ = [str(repo_root / "rag")]
    monkeypatch.setitem(sys.modules, "rag", rag_pkg)
    rag_prompts_pkg = ModuleType("rag.prompts")
    rag_prompts_pkg.__path__ = [str(repo_root / "rag" / "prompts")]
    monkeypatch.setitem(sys.modules, "rag.prompts", rag_prompts_pkg)
    rag_prompts = ModuleType("rag.prompts.generator")
    rag_prompts.chunks_format = lambda *_a, **_kw: ""
    monkeypatch.setitem(sys.modules, "rag.prompts.generator", rag_prompts)

    rag_utils_pkg = ModuleType("rag.utils")
    rag_utils_pkg.__path__ = [str(repo_root / "rag" / "utils")]
    monkeypatch.setitem(sys.modules, "rag.utils", rag_utils_pkg)
    rag_redis = ModuleType("rag.utils.redis_conn")
    rag_redis.REDIS_CONN = SimpleNamespace(delete=lambda *_a, **_kw: None, set=lambda *_a, **_kw: None)
    monkeypatch.setitem(sys.modules, "rag.utils.redis_conn", rag_redis)

    agent_pkg = ModuleType("agent")
    agent_pkg.__path__ = [str(repo_root / "agent")]
    monkeypatch.setitem(sys.modules, "agent", agent_pkg)

    agent_settings = ModuleType("agent.settings")
    agent_settings.FLOAT_ZERO = 1e-8
    agent_settings.PARAM_MAXDEPTH = 5
    monkeypatch.setitem(sys.modules, "agent.settings", agent_settings)

    dsl_migration = ModuleType("agent.dsl_migration")
    dsl_migration.normalize_chunker_dsl = lambda dsl: dsl
    monkeypatch.setitem(sys.modules, "agent.dsl_migration", dsl_migration)

    component_pkg = ModuleType("agent.component")
    component_pkg.__path__ = [str(repo_root / "agent" / "component")]
    monkeypatch.setitem(sys.modules, "agent.component", component_pkg)

    base_spec = importlib.util.spec_from_file_location(
        "agent.component.base", repo_root / "agent" / "component" / "base.py"
    )
    base_mod = importlib.util.module_from_spec(base_spec)
    monkeypatch.setitem(sys.modules, "agent.component.base", base_mod)
    base_spec.loader.exec_module(base_mod)

    iteration_spec = importlib.util.spec_from_file_location(
        "agent.component.iteration", repo_root / "agent" / "component" / "iteration.py"
    )
    iteration_mod = importlib.util.module_from_spec(iteration_spec)
    monkeypatch.setitem(sys.modules, "agent.component.iteration", iteration_mod)
    iteration_spec.loader.exec_module(iteration_mod)

    iterationitem_spec = importlib.util.spec_from_file_location(
        "agent.component.iterationitem",
        repo_root / "agent" / "component" / "iterationitem.py",
    )
    iterationitem_mod = importlib.util.module_from_spec(iterationitem_spec)
    monkeypatch.setitem(sys.modules, "agent.component.iterationitem", iterationitem_mod)
    iterationitem_spec.loader.exec_module(iterationitem_mod)

    class BeginParam(base_mod.ComponentParamBase):
        def check(self):
            return True

    class Begin(base_mod.ComponentBase):
        component_name = "Begin"

        def _invoke(self, **kwargs):
            return

        def thoughts(self):
            return "begin"

    class ProbeParam(base_mod.ComponentParamBase):
        def __init__(self):
            super().__init__()
            self.query = ""
            self.inputs = {"query": {"value": None}}

        def get_input_form(self):
            return {"query": {"name": "Query", "type": "line"}}

        def check(self):
            return True

    class Probe(base_mod.ComponentBase):
        component_name = "Probe"

        def _invoke(self, **kwargs):
            query_text = kwargs.get("query")
            vars_map = self.get_input_elements_from_text(query_text)
            query = self.string_format(
                query_text, {key: value["value"] for key, value in vars_map.items()}
            )
            calls = self._canvas.globals.setdefault("probe.calls", [])
            calls.append(query)
            self.set_output("result", query)

        def thoughts(self):
            return "probe"

    class SinkParam(base_mod.ComponentParamBase):
        def check(self):
            return True

    class Sink(base_mod.ComponentBase):
        component_name = "Sink"

        def _invoke(self, **kwargs):
            self.set_output("done", True)

        def thoughts(self):
            return "sink"

    class UserFillUpParam(base_mod.ComponentParamBase):
        def __init__(self):
            super().__init__()
            self.enable_tips = True
            self.tips = "Please fill"
            self.inputs = {"value": {"type": "line", "name": "Value"}}

        def get_input_form(self):
            return self.inputs

        def check(self):
            return True

    class UserFillUp(base_mod.ComponentBase):
        component_name = "UserFillUp"

        def _invoke(self, **kwargs):
            incoming = kwargs.get("inputs", {})
            if "value" in incoming:
                raw = incoming["value"]
                value = raw.get("value") if isinstance(raw, dict) else raw
                self.set_output("value", value)
            if self._param.enable_tips:
                self.set_output("tips", self._param.tips)

        def get_input_elements(self):
            return self._param.inputs

        def thoughts(self):
            return "fill"

    class MessageParam(base_mod.ComponentParamBase):
        def __init__(self):
            super().__init__()
            self.content = "{UserFillUp:1@value}"

        def check(self):
            return True

    class Message(base_mod.ComponentBase):
        component_name = "Message"

        def _invoke(self, **kwargs):
            self.set_output("content", self.string_format(self._param.content, {}))

        def thoughts(self):
            return "message"

    class_map = {
        "Begin": Begin,
        "BeginParam": BeginParam,
        "Iteration": iteration_mod.Iteration,
        "IterationParam": iteration_mod.IterationParam,
        "IterationItem": iterationitem_mod.IterationItem,
        "IterationItemParam": iterationitem_mod.IterationItemParam,
        "Probe": Probe,
        "ProbeParam": ProbeParam,
        "Sink": Sink,
        "SinkParam": SinkParam,
        "UserFillUp": UserFillUp,
        "UserFillUpParam": UserFillUpParam,
        "Message": Message,
        "MessageParam": MessageParam,
    }

    component_pkg.component_class = lambda name: class_map[name]

    canvas_spec = importlib.util.spec_from_file_location(
        "agent.canvas", repo_root / "agent" / "canvas.py"
    )
    canvas_mod = importlib.util.module_from_spec(canvas_spec)
    monkeypatch.setitem(sys.modules, "agent.canvas", canvas_mod)
    canvas_spec.loader.exec_module(canvas_mod)

    return canvas_mod


async def _collect_events(stream):
    events = []
    async for event in stream:
        events.append(event)
    return events


@pytest.mark.p2
def test_iteration_runtime_processes_all_array_items(monkeypatch):
    canvas_mod = _load_canvas_runtime(monkeypatch)

    dsl = {
        "components": {
            "begin": {
                "obj": {"component_name": "Begin", "params": {}},
                "downstream": ["Iteration:1"],
                "upstream": [],
            },
            "Iteration:1": {
                "obj": {
                    "component_name": "Iteration",
                    "params": {"items_ref": "env.items"},
                },
                "downstream": ["Sink:1"],
                "upstream": ["begin"],
            },
            "IterationItem:1": {
                "obj": {"component_name": "IterationItem", "params": {}},
                "parent_id": "Iteration:1",
                "downstream": ["Probe:1"],
                "upstream": [],
            },
            "Probe:1": {
                "obj": {
                    "component_name": "Probe",
                    "params": {"query": "IterationItem:1@result"},
                },
                "parent_id": "Iteration:1",
                "downstream": [],
                "upstream": ["IterationItem:1"],
            },
            "Sink:1": {
                "obj": {"component_name": "Sink", "params": {}},
                "downstream": [],
                "upstream": ["Iteration:1"],
            },
        },
        "graph": {
            "nodes": [
                {"id": "begin", "data": {"name": "Begin"}},
                {"id": "Iteration:1", "data": {"name": "Iteration"}},
                {"id": "IterationItem:1", "data": {"name": "IterationItem"}},
                {"id": "Probe:1", "data": {"name": "Probe"}},
                {"id": "Sink:1", "data": {"name": "Sink"}},
            ]
        },
        "history": [],
        "path": [],
        "retrieval": [],
        "globals": {
            "sys.query": "",
            "sys.user_id": "",
            "sys.conversation_turns": 0,
            "sys.files": [],
            "sys.history": [],
            "sys.date": "",
            "env.items": ["a", "b", "c"],
        },
    }

    canvas = canvas_mod.Canvas(json.dumps(dsl))
    events = asyncio.run(_collect_events(canvas.run()))

    assert canvas.globals["probe.calls"] == ["a", "b", "c"]
    assert any(event["event"] == "workflow_finished" for event in events)


@pytest.mark.parametrize(
    ("query", "expected_calls"),
    [
        ("{item}", ["a", "b", "c"]),
        ("{index}", ["0", "1", "2"]),
        ("{result}", ["a", "b", "c"]),
    ],
)
@pytest.mark.p2
def test_iteration_runtime_supports_bare_iteration_aliases(monkeypatch, query, expected_calls):
    canvas_mod = _load_canvas_runtime(monkeypatch)

    dsl = {
        "components": {
            "begin": {
                "obj": {"component_name": "Begin", "params": {}},
                "downstream": ["Iteration:1"],
                "upstream": [],
            },
            "Iteration:1": {
                "obj": {
                    "component_name": "Iteration",
                    "params": {"items_ref": "env.items"},
                },
                "downstream": ["Sink:1"],
                "upstream": ["begin"],
            },
            "IterationItem:1": {
                "obj": {"component_name": "IterationItem", "params": {}},
                "parent_id": "Iteration:1",
                "downstream": ["Probe:1"],
                "upstream": [],
            },
            "Probe:1": {
                "obj": {
                    "component_name": "Probe",
                    "params": {"query": query},
                },
                "parent_id": "Iteration:1",
                "downstream": [],
                "upstream": ["IterationItem:1"],
            },
            "Sink:1": {
                "obj": {"component_name": "Sink", "params": {}},
                "downstream": [],
                "upstream": ["Iteration:1"],
            },
        },
        "graph": {
            "nodes": [
                {"id": "begin", "data": {"name": "Begin"}},
                {"id": "Iteration:1", "data": {"name": "Iteration"}},
                {"id": "IterationItem:1", "data": {"name": "IterationItem"}},
                {"id": "Probe:1", "data": {"name": "Probe"}},
                {"id": "Sink:1", "data": {"name": "Sink"}},
            ]
        },
        "history": [],
        "path": [],
        "retrieval": [],
        "globals": {
            "sys.query": "",
            "sys.user_id": "",
            "sys.conversation_turns": 0,
            "sys.files": [],
            "sys.history": [],
            "sys.date": "",
            "env.items": ["a", "b", "c"],
        },
    }

    canvas = canvas_mod.Canvas(json.dumps(dsl))
    asyncio.run(_collect_events(canvas.run()))

    assert canvas.globals["probe.calls"] == expected_calls


@pytest.mark.p2
def test_canvas_resume_does_not_emit_duplicate_workflow_started(monkeypatch):
    canvas_mod = _load_canvas_runtime(monkeypatch)

    dsl = {
        "components": {
            "begin": {
                "obj": {"component_name": "Begin", "params": {}},
                "downstream": ["UserFillUp:1"],
                "upstream": [],
            },
            "UserFillUp:1": {
                "obj": {
                    "component_name": "UserFillUp",
                    "params": {
                        "enable_tips": True,
                        "tips": "Enter value",
                        "inputs": {"value": {"type": "line", "name": "Value"}},
                    },
                },
                "downstream": ["Message:1"],
                "upstream": ["begin"],
            },
            "Message:1": {
                "obj": {
                    "component_name": "Message",
                    "params": {"content": "{UserFillUp:1@value}"},
                },
                "downstream": [],
                "upstream": ["UserFillUp:1"],
            },
        },
        "graph": {
            "nodes": [
                {"id": "begin", "data": {"name": "Begin"}},
                {"id": "UserFillUp:1", "data": {"name": "UserFillUp"}},
                {"id": "Message:1", "data": {"name": "Message"}},
            ]
        },
        "history": [],
        "path": [],
        "retrieval": [],
        "globals": {
            "sys.query": "",
            "sys.user_id": "",
            "sys.conversation_turns": 0,
            "sys.files": [],
            "sys.history": [],
            "sys.date": "",
        },
    }

    canvas = canvas_mod.Canvas(json.dumps(dsl))
    first_events = asyncio.run(_collect_events(canvas.run()))
    first_kinds = [event["event"] for event in first_events]
    assert first_kinds == [
        "workflow_started",
        "node_started",
        "node_finished",
        "user_inputs",
    ]

    resumed_events = asyncio.run(
        _collect_events(canvas.run(query="hello", inputs={"value": {"value": "hello"}}))
    )
    resumed_kinds = [event["event"] for event in resumed_events]
    assert resumed_kinds[0] == "node_started"
    assert "workflow_started" not in resumed_kinds
    assert "message" in resumed_kinds
    assert resumed_kinds[-1] == "workflow_finished"
