import types
from collections.abc import Generator

import pytest

from contexts.wrapper import RecyclableContextVar
from core.datasource.datasource_manager import DatasourceManager
from core.datasource.entities.datasource_entities import DatasourceMessage, DatasourceProviderType
from core.datasource.errors import DatasourceProviderNotFoundError
from dify_graph.entities.workflow_node_execution import WorkflowNodeExecutionStatus
from dify_graph.file import File
from dify_graph.file.enums import FileTransferMethod, FileType
from dify_graph.node_events import StreamChunkEvent, StreamCompletedEvent


def _gen_messages_text_only(text: str) -> Generator[DatasourceMessage, None, None]:
    yield DatasourceMessage(
        type=DatasourceMessage.MessageType.TEXT,
        message=DatasourceMessage.TextMessage(text=text),
        meta=None,
    )


def _drain_generator(gen: Generator[DatasourceMessage, None, object]) -> tuple[list[DatasourceMessage], object | None]:
    messages: list[DatasourceMessage] = []
    try:
        while True:
            messages.append(next(gen))
    except StopIteration as e:
        return messages, e.value


def _invalidate_recyclable_contextvars() -> None:
    """
    Ensure RecyclableContextVar.get() raises LookupError until reset by code under test.
    """
    RecyclableContextVar.increment_thread_recycles()


def test_get_icon_url_calls_runtime(mocker):
    fake_runtime = mocker.Mock()
    fake_runtime.get_icon_url.return_value = "https://icon"
    mocker.patch.object(DatasourceManager, "get_datasource_runtime", return_value=fake_runtime)

    url = DatasourceManager.get_icon_url(
        provider_id="p/x",
        tenant_id="t1",
        datasource_name="ds",
        datasource_type="online_document",
    )
    assert url == "https://icon"
    DatasourceManager.get_datasource_runtime.assert_called_once()


def test_get_datasource_runtime_delegates_to_provider_controller(mocker):
    provider_controller = mocker.Mock()
    provider_controller.get_datasource.return_value = object()
    mocker.patch.object(DatasourceManager, "get_datasource_plugin_provider", return_value=provider_controller)

    runtime = DatasourceManager.get_datasource_runtime(
        provider_id="prov/x",
        datasource_name="ds",
        tenant_id="t1",
        datasource_type=DatasourceProviderType.ONLINE_DOCUMENT,
    )
    assert runtime is provider_controller.get_datasource.return_value
    provider_controller.get_datasource.assert_called_once_with("ds")


@pytest.mark.parametrize(
    ("datasource_type", "controller_path"),
    [
        (
            DatasourceProviderType.ONLINE_DOCUMENT,
            "core.datasource.datasource_manager.OnlineDocumentDatasourcePluginProviderController",
        ),
        (
            DatasourceProviderType.ONLINE_DRIVE,
            "core.datasource.datasource_manager.OnlineDriveDatasourcePluginProviderController",
        ),
        (
            DatasourceProviderType.WEBSITE_CRAWL,
            "core.datasource.datasource_manager.WebsiteCrawlDatasourcePluginProviderController",
        ),
        (
            DatasourceProviderType.LOCAL_FILE,
            "core.datasource.datasource_manager.LocalFileDatasourcePluginProviderController",
        ),
    ],
)
def test_get_datasource_plugin_provider_creates_controller_and_caches(mocker, datasource_type, controller_path):
    _invalidate_recyclable_contextvars()

    provider_entity = types.SimpleNamespace(declaration=object(), plugin_id="plugin", plugin_unique_identifier="uniq")
    fetch = mocker.patch(
        "core.datasource.datasource_manager.PluginDatasourceManager.fetch_datasource_provider",
        return_value=provider_entity,
    )
    ctrl_cls = mocker.patch(controller_path)

    first = DatasourceManager.get_datasource_plugin_provider(
        provider_id=f"prov/{datasource_type.value}",
        tenant_id="t1",
        datasource_type=datasource_type,
    )
    second = DatasourceManager.get_datasource_plugin_provider(
        provider_id=f"prov/{datasource_type.value}",
        tenant_id="t1",
        datasource_type=datasource_type,
    )

    assert first is second
    assert fetch.call_count == 1
    assert ctrl_cls.call_count == 1


def test_get_datasource_plugin_provider_raises_when_provider_entity_missing(mocker):
    _invalidate_recyclable_contextvars()
    mocker.patch(
        "core.datasource.datasource_manager.PluginDatasourceManager.fetch_datasource_provider",
        return_value=None,
    )

    with pytest.raises(DatasourceProviderNotFoundError, match="plugin provider prov/notfound not found"):
        DatasourceManager.get_datasource_plugin_provider(
            provider_id="prov/notfound",
            tenant_id="t1",
            datasource_type=DatasourceProviderType.ONLINE_DOCUMENT,
        )


def test_get_datasource_plugin_provider_raises_for_unsupported_type(mocker):
    _invalidate_recyclable_contextvars()
    provider_entity = types.SimpleNamespace(declaration=object(), plugin_id="plugin", plugin_unique_identifier="uniq")
    mocker.patch(
        "core.datasource.datasource_manager.PluginDatasourceManager.fetch_datasource_provider",
        return_value=provider_entity,
    )

    with pytest.raises(ValueError, match="Unsupported datasource type"):
        DatasourceManager.get_datasource_plugin_provider(
            provider_id="prov/x",
            tenant_id="t1",
            datasource_type=types.SimpleNamespace(),  # not a DatasourceProviderType at runtime
        )


def test_get_datasource_plugin_provider_raises_when_controller_none(mocker):
    _invalidate_recyclable_contextvars()
    provider_entity = types.SimpleNamespace(declaration=object(), plugin_id="plugin", plugin_unique_identifier="uniq")
    mocker.patch(
        "core.datasource.datasource_manager.PluginDatasourceManager.fetch_datasource_provider",
        return_value=provider_entity,
    )
    mocker.patch(
        "core.datasource.datasource_manager.OnlineDocumentDatasourcePluginProviderController",
        return_value=None,
    )

    with pytest.raises(DatasourceProviderNotFoundError, match="Datasource provider prov/x not found"):
        DatasourceManager.get_datasource_plugin_provider(
            provider_id="prov/x",
            tenant_id="t1",
            datasource_type=DatasourceProviderType.ONLINE_DOCUMENT,
        )


def test_stream_online_results_yields_messages_online_document(mocker):
    # stub runtime to yield a text message
    def _doc_messages(**_):
        yield from _gen_messages_text_only("hello")

    fake_runtime = mocker.Mock()
    fake_runtime.get_online_document_page_content.side_effect = _doc_messages
    mocker.patch.object(DatasourceManager, "get_datasource_runtime", return_value=fake_runtime)
    mocker.patch(
        "core.datasource.datasource_manager.DatasourceProviderService.get_datasource_credentials",
        return_value=None,
    )

    gen = DatasourceManager.stream_online_results(
        user_id="u1",
        datasource_name="ds",
        datasource_type="online_document",
        provider_id="p/x",
        tenant_id="t1",
        provider="prov",
        plugin_id="plug",
        credential_id="",
        datasource_param=types.SimpleNamespace(workspace_id="w", page_id="pg", type="t"),
        online_drive_request=None,
    )
    msgs = list(gen)
    assert len(msgs) == 1
    assert msgs[0].message.text == "hello"


def test_stream_online_results_sets_credentials_and_returns_empty_dict_online_document(mocker):
    class _Runtime:
        def __init__(self) -> None:
            self.runtime = types.SimpleNamespace(credentials=None)

        def get_online_document_page_content(self, **_kwargs):
            yield from _gen_messages_text_only("hello")

    runtime = _Runtime()
    mocker.patch.object(DatasourceManager, "get_datasource_runtime", return_value=runtime)
    mocker.patch(
        "core.datasource.datasource_manager.DatasourceProviderService.get_datasource_credentials",
        return_value={"token": "t"},
    )

    gen = DatasourceManager.stream_online_results(
        user_id="u1",
        datasource_name="ds",
        datasource_type="online_document",
        provider_id="p/x",
        tenant_id="t1",
        provider="prov",
        plugin_id="plug",
        credential_id="cred",
        datasource_param=types.SimpleNamespace(workspace_id="w", page_id="pg", type="t"),
        online_drive_request=None,
    )
    messages, final_value = _drain_generator(gen)

    assert runtime.runtime.credentials == {"token": "t"}
    assert [m.message.text for m in messages] == ["hello"]
    assert final_value == {}


def test_stream_online_results_raises_when_missing_params(mocker):
    class _Runtime:
        def __init__(self) -> None:
            self.runtime = types.SimpleNamespace(credentials=None)

        def get_online_document_page_content(self, **_kwargs):
            yield from _gen_messages_text_only("never")

        def online_drive_download_file(self, **_kwargs):
            yield from _gen_messages_text_only("never")

    mocker.patch.object(DatasourceManager, "get_datasource_runtime", return_value=_Runtime())
    mocker.patch(
        "core.datasource.datasource_manager.DatasourceProviderService.get_datasource_credentials",
        return_value={},
    )

    with pytest.raises(ValueError, match="datasource_param is required for ONLINE_DOCUMENT streaming"):
        list(
            DatasourceManager.stream_online_results(
                user_id="u1",
                datasource_name="ds",
                datasource_type="online_document",
                provider_id="p/x",
                tenant_id="t1",
                provider="prov",
                plugin_id="plug",
                credential_id="",
                datasource_param=None,
                online_drive_request=None,
            )
        )

    with pytest.raises(ValueError, match="online_drive_request is required for ONLINE_DRIVE streaming"):
        list(
            DatasourceManager.stream_online_results(
                user_id="u1",
                datasource_name="ds",
                datasource_type="online_drive",
                provider_id="p/x",
                tenant_id="t1",
                provider="prov",
                plugin_id="plug",
                credential_id="",
                datasource_param=None,
                online_drive_request=None,
            )
        )


def test_stream_online_results_yields_messages_and_returns_empty_dict_online_drive(mocker):
    class _Runtime:
        def __init__(self) -> None:
            self.runtime = types.SimpleNamespace(credentials=None)

        def online_drive_download_file(self, **_kwargs):
            yield from _gen_messages_text_only("drive")

    runtime = _Runtime()
    mocker.patch.object(DatasourceManager, "get_datasource_runtime", return_value=runtime)
    mocker.patch(
        "core.datasource.datasource_manager.DatasourceProviderService.get_datasource_credentials",
        return_value={"token": "t"},
    )

    gen = DatasourceManager.stream_online_results(
        user_id="u1",
        datasource_name="ds",
        datasource_type="online_drive",
        provider_id="p/x",
        tenant_id="t1",
        provider="prov",
        plugin_id="plug",
        credential_id="cred",
        datasource_param=None,
        online_drive_request=types.SimpleNamespace(id="fid", bucket="b"),
    )
    messages, final_value = _drain_generator(gen)

    assert runtime.runtime.credentials == {"token": "t"}
    assert [m.message.text for m in messages] == ["drive"]
    assert final_value == {}


def test_stream_online_results_raises_for_unsupported_stream_type(mocker):
    mocker.patch.object(DatasourceManager, "get_datasource_runtime", return_value=mocker.Mock())
    mocker.patch(
        "core.datasource.datasource_manager.DatasourceProviderService.get_datasource_credentials",
        return_value={},
    )

    with pytest.raises(ValueError, match="Unsupported datasource type for streaming"):
        list(
            DatasourceManager.stream_online_results(
                user_id="u1",
                datasource_name="ds",
                datasource_type="website_crawl",
                provider_id="p/x",
                tenant_id="t1",
                provider="prov",
                plugin_id="plug",
                credential_id="",
                datasource_param=None,
                online_drive_request=None,
            )
        )


def test_stream_node_events_emits_events_online_document(mocker):
    # make manager's low-level stream produce TEXT only
    mocker.patch.object(
        DatasourceManager,
        "stream_online_results",
        return_value=_gen_messages_text_only("hello"),
    )

    events = list(
        DatasourceManager.stream_node_events(
            node_id="nodeA",
            user_id="u1",
            datasource_name="ds",
            datasource_type="online_document",
            provider_id="p/x",
            tenant_id="t1",
            provider="prov",
            plugin_id="plug",
            credential_id="",
            parameters_for_log={"k": "v"},
            datasource_info={"user_id": "u1"},
            variable_pool=mocker.Mock(),
            datasource_param=types.SimpleNamespace(workspace_id="w", page_id="pg", type="t"),
            online_drive_request=None,
        )
    )
    # should contain one StreamChunkEvent then a final chunk (empty) and a completed event
    assert isinstance(events[0], StreamChunkEvent)
    assert events[0].chunk == "hello"
    assert isinstance(events[-1], StreamCompletedEvent)
    assert events[-1].node_run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED


def test_stream_node_events_builds_file_and_variables_from_messages(mocker):
    mocker.patch.object(DatasourceManager, "stream_online_results", return_value=_gen_messages_text_only("ignored"))

    def _transformed(**_kwargs):
        yield DatasourceMessage(
            type=DatasourceMessage.MessageType.IMAGE_LINK,
            message=DatasourceMessage.TextMessage(text="/files/datasources/tool_file_1.png"),
            meta={},
        )
        yield DatasourceMessage(
            type=DatasourceMessage.MessageType.TEXT,
            message=DatasourceMessage.TextMessage(text="hello"),
            meta=None,
        )
        yield DatasourceMessage(
            type=DatasourceMessage.MessageType.LINK,
            message=DatasourceMessage.TextMessage(text="http://example.com"),
            meta=None,
        )
        yield DatasourceMessage(
            type=DatasourceMessage.MessageType.VARIABLE,
            message=DatasourceMessage.VariableMessage(variable_name="v", variable_value="a", stream=True),
            meta=None,
        )
        yield DatasourceMessage(
            type=DatasourceMessage.MessageType.VARIABLE,
            message=DatasourceMessage.VariableMessage(variable_name="v", variable_value="b", stream=True),
            meta=None,
        )
        yield DatasourceMessage(
            type=DatasourceMessage.MessageType.VARIABLE,
            message=DatasourceMessage.VariableMessage(variable_name="x", variable_value=1, stream=False),
            meta=None,
        )
        yield DatasourceMessage(
            type=DatasourceMessage.MessageType.JSON,
            message=DatasourceMessage.JsonMessage(json_object={"k": "v"}),
            meta=None,
        )

    mocker.patch(
        "core.datasource.datasource_manager.DatasourceFileMessageTransformer.transform_datasource_invoke_messages",
        side_effect=_transformed,
    )

    fake_tool_file = types.SimpleNamespace(mimetype="image/png")

    class _Session:
        def __enter__(self):
            return self

        def __exit__(self, *exc):
            return False

        def scalar(self, _stmt):
            return fake_tool_file

    mocker.patch("core.datasource.datasource_manager.session_factory.create_session", return_value=_Session())
    mocker.patch(
        "core.datasource.datasource_manager.file_factory.get_file_type_by_mime_type", return_value=FileType.IMAGE
    )
    built = File(
        tenant_id="t1",
        type=FileType.IMAGE,
        transfer_method=FileTransferMethod.TOOL_FILE,
        related_id="tool_file_1",
        extension=".png",
        mime_type="image/png",
        storage_key="k",
    )
    build_from_mapping = mocker.patch(
        "core.datasource.datasource_manager.file_factory.build_from_mapping",
        return_value=built,
    )

    variable_pool = mocker.Mock()

    events = list(
        DatasourceManager.stream_node_events(
            node_id="nodeA",
            user_id="u1",
            datasource_name="ds",
            datasource_type="online_document",
            provider_id="p/x",
            tenant_id="t1",
            provider="prov",
            plugin_id="plug",
            credential_id="",
            parameters_for_log={"k": "v"},
            datasource_info={"info": "x"},
            variable_pool=variable_pool,
            datasource_param=types.SimpleNamespace(workspace_id="w", page_id="pg", type="t"),
            online_drive_request=None,
        )
    )

    build_from_mapping.assert_called_once()
    variable_pool.add.assert_not_called()

    assert any(isinstance(e, StreamChunkEvent) and e.chunk == "hello" for e in events)
    assert any(isinstance(e, StreamChunkEvent) and e.chunk.startswith("Link: http") for e in events)
    assert any(isinstance(e, StreamChunkEvent) and e.selector == ["nodeA", "v"] and e.chunk == "a" for e in events)
    assert any(isinstance(e, StreamChunkEvent) and e.selector == ["nodeA", "v"] and e.chunk == "b" for e in events)
    assert isinstance(events[-2], StreamChunkEvent)
    assert events[-2].is_final is True

    assert isinstance(events[-1], StreamCompletedEvent)
    assert events[-1].node_run_result.outputs["v"] == "ab"
    assert events[-1].node_run_result.outputs["x"] == 1


def test_stream_node_events_raises_when_toolfile_missing(mocker):
    mocker.patch.object(DatasourceManager, "stream_online_results", return_value=_gen_messages_text_only("ignored"))

    def _transformed(**_kwargs):
        yield DatasourceMessage(
            type=DatasourceMessage.MessageType.IMAGE_LINK,
            message=DatasourceMessage.TextMessage(text="/files/datasources/missing.png"),
            meta={},
        )

    mocker.patch(
        "core.datasource.datasource_manager.DatasourceFileMessageTransformer.transform_datasource_invoke_messages",
        side_effect=_transformed,
    )

    class _Session:
        def __enter__(self):
            return self

        def __exit__(self, *exc):
            return False

        def scalar(self, _stmt):
            return None

    mocker.patch("core.datasource.datasource_manager.session_factory.create_session", return_value=_Session())

    with pytest.raises(ValueError, match="ToolFile not found for file_id=missing, tenant_id=t1"):
        list(
            DatasourceManager.stream_node_events(
                node_id="nodeA",
                user_id="u1",
                datasource_name="ds",
                datasource_type="online_document",
                provider_id="p/x",
                tenant_id="t1",
                provider="prov",
                plugin_id="plug",
                credential_id="",
                parameters_for_log={},
                datasource_info={},
                variable_pool=mocker.Mock(),
                datasource_param=types.SimpleNamespace(workspace_id="w", page_id="pg", type="t"),
                online_drive_request=None,
            )
        )


def test_stream_node_events_online_drive_sets_variable_pool_file_and_outputs(mocker):
    mocker.patch.object(DatasourceManager, "stream_online_results", return_value=_gen_messages_text_only("ignored"))

    file_in = File(
        tenant_id="t1",
        type=FileType.DOCUMENT,
        transfer_method=FileTransferMethod.TOOL_FILE,
        related_id="tf",
        extension=".pdf",
        mime_type="application/pdf",
        storage_key="k",
    )

    def _transformed(**_kwargs):
        yield DatasourceMessage(
            type=DatasourceMessage.MessageType.FILE,
            message=DatasourceMessage.FileMessage(file_marker="file_marker"),
            meta={"file": file_in},
        )

    mocker.patch(
        "core.datasource.datasource_manager.DatasourceFileMessageTransformer.transform_datasource_invoke_messages",
        side_effect=_transformed,
    )

    variable_pool = mocker.Mock()
    events = list(
        DatasourceManager.stream_node_events(
            node_id="nodeA",
            user_id="u1",
            datasource_name="ds",
            datasource_type="online_drive",
            provider_id="p/x",
            tenant_id="t1",
            provider="prov",
            plugin_id="plug",
            credential_id="",
            parameters_for_log={},
            datasource_info={"k": "v"},
            variable_pool=variable_pool,
            datasource_param=None,
            online_drive_request=types.SimpleNamespace(id="id", bucket="b"),
        )
    )

    variable_pool.add.assert_called_once()
    assert variable_pool.add.call_args[0][0] == ["nodeA", "file"]
    assert variable_pool.add.call_args[0][1] == file_in

    completed = events[-1]
    assert isinstance(completed, StreamCompletedEvent)
    assert completed.node_run_result.outputs["file"] == file_in
    assert completed.node_run_result.outputs["datasource_type"] == DatasourceProviderType.ONLINE_DRIVE


def test_stream_node_events_skips_file_build_for_non_online_types(mocker):
    mocker.patch.object(DatasourceManager, "stream_online_results", return_value=_gen_messages_text_only("ignored"))

    def _transformed(**_kwargs):
        yield DatasourceMessage(
            type=DatasourceMessage.MessageType.IMAGE_LINK,
            message=DatasourceMessage.TextMessage(text="/files/datasources/tool_file_1.png"),
            meta={},
        )

    mocker.patch(
        "core.datasource.datasource_manager.DatasourceFileMessageTransformer.transform_datasource_invoke_messages",
        side_effect=_transformed,
    )
    build_from_mapping = mocker.patch("core.datasource.datasource_manager.file_factory.build_from_mapping")

    events = list(
        DatasourceManager.stream_node_events(
            node_id="nodeA",
            user_id="u1",
            datasource_name="ds",
            datasource_type="website_crawl",
            provider_id="p/x",
            tenant_id="t1",
            provider="prov",
            plugin_id="plug",
            credential_id="",
            parameters_for_log={},
            datasource_info={},
            variable_pool=mocker.Mock(),
            datasource_param=None,
            online_drive_request=None,
        )
    )

    build_from_mapping.assert_not_called()
    assert isinstance(events[-1], StreamCompletedEvent)
    assert events[-1].node_run_result.outputs["file"] is None


def test_get_upload_file_by_id_builds_file(mocker):
    # fake UploadFile row
    fake_row = types.SimpleNamespace(
        id="fid",
        name="f",
        extension="txt",
        mime_type="text/plain",
        size=1,
        key="k",
        source_url="http://x",
    )

    class _Q:
        def __init__(self, row):
            self._row = row

        def where(self, *_args, **_kwargs):
            return self

        def first(self):
            return self._row

    class _S:
        def __init__(self, row):
            self._row = row

        def __enter__(self):
            return self

        def __exit__(self, *exc):
            return False

        def query(self, *_):
            return _Q(self._row)

    mocker.patch("core.datasource.datasource_manager.session_factory.create_session", return_value=_S(fake_row))

    f = DatasourceManager.get_upload_file_by_id(file_id="fid", tenant_id="t1")
    assert f.related_id == "fid"
    assert f.extension == ".txt"


def test_get_upload_file_by_id_raises_when_missing(mocker):
    class _Q:
        def where(self, *_args, **_kwargs):
            return self

        def first(self):
            return None

    class _S:
        def __enter__(self):
            return self

        def __exit__(self, *exc):
            return False

        def query(self, *_):
            return _Q()

    mocker.patch("core.datasource.datasource_manager.session_factory.create_session", return_value=_S())

    with pytest.raises(ValueError, match="UploadFile not found for file_id=fid, tenant_id=t1"):
        DatasourceManager.get_upload_file_by_id(file_id="fid", tenant_id="t1")
