from contextlib import contextmanager
from types import SimpleNamespace
from unittest.mock import ANY, MagicMock, Mock, patch

import pytest

from core.app.apps.chat.app_generator import ChatAppGenerator
from core.app.apps.chat.app_runner import ChatAppRunner
from core.app.apps.exc import GenerateTaskStoppedError
from core.app.entities.app_invoke_entities import InvokeFrom
from core.app.entities.queue_entities import QueueAnnotationReplyEvent
from core.moderation.base import ModerationError
from graphon.model_runtime.errors.invoke import InvokeAuthorizationError
from models.model import AppMode


class DummyGenerateEntity:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)


class DummyQueueManager:
    def __init__(self, *args, **kwargs):
        self.published = []

    def publish_error(self, error, pub_from):
        self.published.append((error, pub_from))

    def publish(self, event, pub_from):
        self.published.append((event, pub_from))


@contextmanager
def patched_create_session(*, return_value=None, side_effect=None):
    session = MagicMock()
    if side_effect is not None:
        session.scalar.side_effect = side_effect
    else:
        session.scalar.return_value = return_value
    session_context = MagicMock()
    session_context.__enter__.return_value = session
    with patch("core.app.apps.chat.app_runner.create_session", return_value=session_context):
        yield session


class TestChatAppGenerator:
    def test_generate_requires_query(self):
        generator = ChatAppGenerator()
        with pytest.raises(ValueError):
            generator.generate(
                session=MagicMock(),
                app_model=SimpleNamespace(),
                user=SimpleNamespace(),
                args={"inputs": {}},
                invoke_from=InvokeFrom.SERVICE_API,
                streaming=False,
            )

    def test_generate_rejects_non_string_query(self):
        generator = ChatAppGenerator()
        with pytest.raises(ValueError):
            generator.generate(
                session=MagicMock(),
                app_model=SimpleNamespace(),
                user=SimpleNamespace(),
                args={"query": 1, "inputs": {}},
                invoke_from=InvokeFrom.SERVICE_API,
                streaming=False,
            )

    def test_generate_debugger_overrides_model_config(self):
        generator = ChatAppGenerator()
        app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1")
        user = SimpleNamespace(id="user-1", session_id="session-1")
        session = MagicMock()
        args = {
            "query": "hi",
            "inputs": {},
            "conversation_id": "conversation-1",
            "model_config": {"foo": "bar"},
            "trace_session_id": "session-1",
        }

        with (
            patch(
                "core.app.apps.chat.app_generator.ConversationService.get_conversation", return_value=None
            ) as get_conversation,
            patch("core.app.apps.chat.app_generator.ChatAppConfigManager.config_validate", return_value={"x": 1}),
            patch(
                "core.app.apps.chat.app_generator.ChatAppConfigManager.get_app_config",
                return_value=SimpleNamespace(
                    variables=[], external_data_variables=[], app_model_config_dict={}, app_mode=AppMode.CHAT
                ),
            ),
            patch("core.app.apps.chat.app_generator.ModelConfigConverter.convert", return_value=SimpleNamespace()),
            patch("core.app.apps.chat.app_generator.FileUploadConfigManager.convert", return_value=None),
            patch("core.app.apps.chat.app_generator.file_factory.build_from_mappings", return_value=[]),
            patch(
                "core.app.apps.chat.app_generator.ChatAppGenerateEntity",
                Mock(side_effect=DummyGenerateEntity),
            ) as generate_entity,
            patch("core.app.apps.chat.app_generator.TraceQueueManager", return_value=SimpleNamespace()),
            patch("core.app.apps.chat.app_generator.MessageBasedAppQueueManager", DummyQueueManager),
            patch(
                "core.app.apps.chat.app_generator.ChatAppGenerateResponseConverter.convert", return_value={"ok": True}
            ),
            patch.object(ChatAppGenerator, "_get_app_model_config", return_value=SimpleNamespace(to_dict=lambda: {})),
            patch.object(ChatAppGenerator, "_prepare_user_inputs", return_value={}),
            patch.object(
                ChatAppGenerator,
                "_init_generate_records",
                return_value=(SimpleNamespace(id="c1", mode="chat"), SimpleNamespace(id="m1")),
            ),
            patch.object(ChatAppGenerator, "_handle_response", return_value={"response": True}),
            patch("core.app.apps.chat.app_generator.copy_current_request_context", side_effect=lambda f: f),
            patch("core.app.apps.chat.app_generator.threading.Thread") as mock_thread,
        ):
            mock_thread.return_value.start.return_value = None
            result = generator.generate(app_model, user, args, InvokeFrom.DEBUGGER, streaming=False, session=session)

        assert result == {"ok": True}
        assert get_conversation.call_args.kwargs["session"] is session
        assert generate_entity.call_args.kwargs["extras"]["trace_session_id"] == "session-1"

    def test_generate_uses_session_for_annotation_reply(self):
        generator = ChatAppGenerator()
        app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1")
        app_model_config = MagicMock(id="config-1", app_id="app-1")
        annotation_reply = {"enabled": False}
        user = SimpleNamespace(id="user-1", session_id="session-1")
        session = MagicMock()

        with (
            patch.object(ChatAppGenerator, "_get_app_model_config", return_value=app_model_config),
            patch(
                "core.app.apps.chat.app_generator.load_annotation_reply_config",
                return_value=annotation_reply,
            ) as load_annotation_reply_config,
            patch("core.app.apps.chat.app_generator.FileUploadConfigManager.convert", return_value=None),
            patch(
                "core.app.apps.chat.app_generator.ChatAppConfigManager.get_app_config",
                side_effect=RuntimeError("stop after app config"),
            ) as get_app_config,
        ):
            with pytest.raises(RuntimeError, match="stop after app config"):
                generator.generate(
                    app_model,
                    user,
                    {"query": "hi", "inputs": {}},
                    InvokeFrom.WEB_APP,
                    session=session,
                )

        load_annotation_reply_config.assert_called_once_with(session, "app-1")
        app_model_config.to_dict.assert_called_once_with(annotation_reply=annotation_reply)
        assert get_app_config.call_args.kwargs["annotation_reply"] is annotation_reply

    def test_generate_rejects_model_config_override_for_non_debugger(self):
        generator = ChatAppGenerator()
        with pytest.raises(ValueError):
            with (
                patch.object(
                    ChatAppGenerator, "_get_app_model_config", return_value=SimpleNamespace(to_dict=lambda: {})
                ),
            ):
                generator.generate(
                    session=MagicMock(),
                    app_model=SimpleNamespace(tenant_id="t1", id="a1", mode=AppMode.CHAT.value),
                    user=SimpleNamespace(id="u1", session_id="s1"),
                    args={"query": "hi", "inputs": {}, "model_config": {"foo": "bar"}},
                    invoke_from=InvokeFrom.SERVICE_API,
                    streaming=False,
                )

    def test_generate_worker_handles_exceptions(self):
        generator = ChatAppGenerator()
        queue_manager = DummyQueueManager()
        entity = DummyGenerateEntity(task_id="t1", user_id="u1")

        with (
            patch.object(ChatAppGenerator, "_get_conversation", return_value=SimpleNamespace()),
            patch.object(ChatAppGenerator, "_get_message", return_value=SimpleNamespace()),
            patch("core.app.apps.chat.app_generator.ChatAppRunner.run", side_effect=InvokeAuthorizationError()),
            patch("core.app.apps.chat.app_generator.session_factory.create_session") as create_session,
            patch("core.app.apps.chat.app_generator.db.session.close"),
        ):
            create_session.return_value.__enter__.return_value = MagicMock()
            generator._generate_worker(
                flask_app=Mock(app_context=Mock(return_value=Mock(__enter__=Mock(), __exit__=Mock()))),
                application_generate_entity=entity,
                queue_manager=queue_manager,
                conversation_id="c1",
                message_id="m1",
            )

        assert queue_manager.published

        with (
            patch.object(ChatAppGenerator, "_get_conversation", return_value=SimpleNamespace()),
            patch.object(ChatAppGenerator, "_get_message", return_value=SimpleNamespace()),
            patch("core.app.apps.chat.app_generator.ChatAppRunner.run", side_effect=GenerateTaskStoppedError()),
            patch("core.app.apps.chat.app_generator.session_factory.create_session") as create_session,
            patch("core.app.apps.chat.app_generator.db.session.close"),
        ):
            create_session.return_value.__enter__.return_value = MagicMock()
            generator._generate_worker(
                flask_app=Mock(app_context=Mock(return_value=Mock(__enter__=Mock(), __exit__=Mock()))),
                application_generate_entity=entity,
                queue_manager=queue_manager,
                conversation_id="c1",
                message_id="m1",
            )


class TestChatAppRunner:
    def test_run_raises_when_app_missing(self):
        runner = ChatAppRunner()
        app_config = SimpleNamespace(
            app_id="app-1", tenant_id="tenant-1", prompt_template=None, external_data_variables=[]
        )
        app_generate_entity = DummyGenerateEntity(
            app_config=app_config,
            model_conf=SimpleNamespace(provider_model_bundle=None, model=None, parameters={}, app_model_config_dict={}),
            inputs={},
            query="hi",
            files=[],
            file_upload_config=None,
            conversation_id=None,
            stream=False,
            user_id="user-1",
            invoke_from=InvokeFrom.SERVICE_API,
        )

        with patched_create_session(return_value=None):
            with pytest.raises(ValueError):
                runner.run(
                    app_generate_entity, DummyQueueManager(), SimpleNamespace(), SimpleNamespace(id="m1"), MagicMock()
                )

    def test_run_moderation_error_direct_output(self):
        runner = ChatAppRunner()
        app_config = SimpleNamespace(
            app_id="app-1",
            tenant_id="tenant-1",
            prompt_template=None,
            external_data_variables=[],
            dataset=None,
            additional_features=None,
        )
        app_generate_entity = DummyGenerateEntity(
            app_config=app_config,
            model_conf=SimpleNamespace(provider_model_bundle=None, model=None, parameters={}, app_model_config_dict={}),
            inputs={},
            query="hi",
            files=[],
            file_upload_config=None,
            conversation_id=None,
            stream=False,
            user_id="user-1",
            invoke_from=InvokeFrom.SERVICE_API,
        )

        with (
            patched_create_session(return_value=SimpleNamespace(id="app-1", tenant_id="tenant-1")),
            patch.object(ChatAppRunner, "organize_prompt_messages", return_value=([], [])),
            patch.object(ChatAppRunner, "moderation_for_inputs", side_effect=ModerationError("blocked")),
            patch.object(ChatAppRunner, "direct_output") as mock_direct,
        ):
            runner.run(
                app_generate_entity, DummyQueueManager(), SimpleNamespace(), SimpleNamespace(id="m1"), MagicMock()
            )

        mock_direct.assert_called_once()

    def test_run_annotation_reply_short_circuits(self):
        runner = ChatAppRunner()
        app_config = SimpleNamespace(
            app_id="app-1",
            tenant_id="tenant-1",
            prompt_template=None,
            external_data_variables=[],
            dataset=None,
            additional_features=None,
        )
        app_generate_entity = DummyGenerateEntity(
            app_config=app_config,
            model_conf=SimpleNamespace(provider_model_bundle=None, model=None, parameters={}, app_model_config_dict={}),
            inputs={},
            query="hi",
            files=[],
            file_upload_config=None,
            conversation_id=None,
            stream=False,
            user_id="user-1",
            invoke_from=InvokeFrom.SERVICE_API,
        )

        annotation = SimpleNamespace(id="ann-1", content="answer")

        with (
            patched_create_session(return_value=SimpleNamespace(id="app-1", tenant_id="tenant-1")),
            patch.object(ChatAppRunner, "organize_prompt_messages", return_value=([], [])),
            patch.object(ChatAppRunner, "moderation_for_inputs", return_value=(None, {}, "hi")),
            patch.object(ChatAppRunner, "query_app_annotations_to_reply", return_value=annotation) as annotation_query,
            patch.object(ChatAppRunner, "direct_output") as mock_direct,
        ):
            queue_manager = DummyQueueManager()
            write_session = MagicMock()
            runner.run(app_generate_entity, queue_manager, SimpleNamespace(), SimpleNamespace(id="m1"), write_session)

        assert any(isinstance(item[0], QueueAnnotationReplyEvent) for item in queue_manager.published)
        assert annotation_query.call_args.kwargs["session"] is write_session
        mock_direct.assert_called_once()

    def test_run_returns_when_hosting_moderation_blocks(self):
        runner = ChatAppRunner()
        app_config = SimpleNamespace(
            app_id="app-1",
            tenant_id="tenant-1",
            prompt_template=None,
            external_data_variables=[],
            dataset=None,
            additional_features=None,
        )
        app_generate_entity = DummyGenerateEntity(
            app_config=app_config,
            model_conf=SimpleNamespace(provider_model_bundle=None, model=None, parameters={}, app_model_config_dict={}),
            inputs={},
            query="hi",
            files=[],
            file_upload_config=None,
            conversation_id=None,
            stream=False,
            user_id="user-1",
            invoke_from=InvokeFrom.SERVICE_API,
        )

        with (
            patched_create_session(return_value=SimpleNamespace(id="app-1", tenant_id="tenant-1")),
            patch.object(ChatAppRunner, "organize_prompt_messages", return_value=([], [])),
            patch.object(ChatAppRunner, "moderation_for_inputs", return_value=(None, {}, "hi")),
            patch.object(ChatAppRunner, "query_app_annotations_to_reply", return_value=None),
            patch.object(ChatAppRunner, "check_hosting_moderation", return_value=True),
        ):
            runner.run(
                app_generate_entity, DummyQueueManager(), SimpleNamespace(), SimpleNamespace(id="m1"), MagicMock()
            )

    def test_run_closes_explicit_session_before_stream_consumption(self):
        runner = ChatAppRunner()
        app_config = SimpleNamespace(
            app_id="app-1",
            tenant_id="tenant-1",
            prompt_template=None,
            external_data_variables=[],
            dataset=None,
            additional_features=None,
        )
        app_generate_entity = DummyGenerateEntity(
            app_config=app_config,
            model_conf=SimpleNamespace(provider_model_bundle=None, model="model-1", parameters={}),
            inputs={},
            query="hi",
            files=[],
            file_upload_config=None,
            conversation_id=None,
            stream=True,
            user_id="user-1",
            invoke_from=InvokeFrom.SERVICE_API,
        )

        events = []
        queue_manager = DummyQueueManager()
        model_instance = MagicMock()
        session = MagicMock()
        session.commit.side_effect = lambda: events.append("commit")
        session.close.side_effect = lambda: events.append("close")

        def invoke_stream():
            events.append("first-chunk")
            yield "chunk"

        def invoke_llm(**kwargs):
            events.append("invoke")
            return invoke_stream()

        with (
            patched_create_session(return_value=SimpleNamespace(id="app-1", tenant_id="tenant-1")),
            patch.object(ChatAppRunner, "organize_prompt_messages", return_value=([], [])),
            patch.object(ChatAppRunner, "moderation_for_inputs", return_value=(None, {}, "hi")),
            patch.object(ChatAppRunner, "query_app_annotations_to_reply", return_value=None),
            patch.object(ChatAppRunner, "check_hosting_moderation", return_value=False),
            patch.object(ChatAppRunner, "recalc_llm_max_tokens"),
            patch.object(
                ChatAppRunner,
                "_handle_invoke_result",
                side_effect=lambda invoke_result, **kwargs: list(invoke_result),
            ) as mock_handle,
            patch("core.app.apps.chat.app_runner.ModelInstance", return_value=model_instance),
        ):
            model_instance.invoke_llm.side_effect = invoke_llm
            runner.run(app_generate_entity, queue_manager, SimpleNamespace(), SimpleNamespace(id="m1"), session)

        assert events == ["commit", "close", "commit", "close", "invoke", "first-chunk"]
        mock_handle.assert_called_once_with(
            invoke_result=ANY,
            queue_manager=queue_manager,
            stream=True,
            message_id="m1",
            user_id="user-1",
            tenant_id="tenant-1",
        )
