from collections.abc import Sequence
from typing import Optional

import pytest

from core.app.entities.app_invoke_entities import InvokeFrom, ModelConfigWithCredentialsEntity
from core.entities.provider_configuration import ProviderConfiguration, ProviderModelBundle
from core.entities.provider_entities import CustomConfiguration, SystemConfiguration
from core.file import File, FileTransferMethod, FileType
from core.model_runtime.entities.common_entities import I18nObject
from core.model_runtime.entities.message_entities import (
    ImagePromptMessageContent,
    PromptMessage,
    PromptMessageRole,
    TextPromptMessageContent,
    UserPromptMessage,
)
from core.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType
from core.model_runtime.model_providers.model_provider_factory import ModelProviderFactory
from core.variables import ArrayAnySegment, ArrayFileSegment, NoneSegment
from core.workflow.entities.variable_pool import VariablePool
from core.workflow.graph_engine import Graph, GraphInitParams, GraphRuntimeState
from core.workflow.nodes.answer import AnswerStreamGenerateRoute
from core.workflow.nodes.end import EndStreamParam
from core.workflow.nodes.llm.entities import (
    ContextConfig,
    LLMNodeChatModelMessage,
    LLMNodeData,
    ModelConfig,
    VisionConfig,
    VisionConfigOptions,
)
from core.workflow.nodes.llm.node import LLMNode
from models.enums import UserFrom
from models.provider import ProviderType
from models.workflow import WorkflowType


class MockTokenBufferMemory:
    def __init__(self, history_messages=None):
        self.history_messages = history_messages or []

    def get_history_prompt_messages(
        self, max_token_limit: int = 2000, message_limit: Optional[int] = None
    ) -> Sequence[PromptMessage]:
        if message_limit is not None:
            return self.history_messages[-message_limit * 2 :]
        return self.history_messages


@pytest.fixture
def llm_node():
    data = LLMNodeData(
        title="Test LLM",
        model=ModelConfig(provider="openai", name="gpt-3.5-turbo", mode="chat", completion_params={}),
        prompt_template=[],
        memory=None,
        context=ContextConfig(enabled=False),
        vision=VisionConfig(
            enabled=True,
            configs=VisionConfigOptions(
                variable_selector=["sys", "files"],
                detail=ImagePromptMessageContent.DETAIL.HIGH,
            ),
        ),
    )
    variable_pool = VariablePool(
        system_variables={},
        user_inputs={},
    )
    node = LLMNode(
        id="1",
        config={
            "id": "1",
            "data": data.model_dump(),
        },
        graph_init_params=GraphInitParams(
            tenant_id="1",
            app_id="1",
            workflow_type=WorkflowType.WORKFLOW,
            workflow_id="1",
            graph_config={},
            user_id="1",
            user_from=UserFrom.ACCOUNT,
            invoke_from=InvokeFrom.SERVICE_API,
            call_depth=0,
        ),
        graph=Graph(
            root_node_id="1",
            answer_stream_generate_routes=AnswerStreamGenerateRoute(
                answer_dependencies={},
                answer_generate_route={},
            ),
            end_stream_param=EndStreamParam(
                end_dependencies={},
                end_stream_variable_selector_mapping={},
            ),
        ),
        graph_runtime_state=GraphRuntimeState(
            variable_pool=variable_pool,
            start_at=0,
        ),
    )
    return node


@pytest.fixture
def model_config():
    # Create actual provider and model type instances
    model_provider_factory = ModelProviderFactory(tenant_id="test")
    provider_instance = model_provider_factory.get_plugin_model_provider("openai")
    model_type_instance = model_provider_factory.get_model_type_instance("openai", ModelType.LLM)

    # Create a ProviderModelBundle
    provider_model_bundle = ProviderModelBundle(
        configuration=ProviderConfiguration(
            tenant_id="1",
            provider=provider_instance,
            preferred_provider_type=ProviderType.CUSTOM,
            using_provider_type=ProviderType.CUSTOM,
            system_configuration=SystemConfiguration(enabled=False),
            custom_configuration=CustomConfiguration(provider=None),
            model_settings=[],
        ),
        model_type_instance=model_type_instance,
    )

    # Create and return a ModelConfigWithCredentialsEntity
    return ModelConfigWithCredentialsEntity(
        provider="openai",
        model="gpt-3.5-turbo",
        model_schema=AIModelEntity(
            model="gpt-3.5-turbo",
            label=I18nObject(en_US="GPT-3.5 Turbo"),
            model_type=ModelType.LLM,
            fetch_from=FetchFrom.CUSTOMIZABLE_MODEL,
            model_properties={},
        ),
        mode="chat",
        credentials={},
        parameters={},
        provider_model_bundle=provider_model_bundle,
    )


def test_fetch_files_with_file_segment(llm_node):
    file = File(
        id="1",
        tenant_id="test",
        type=FileType.IMAGE,
        filename="test.jpg",
        transfer_method=FileTransferMethod.LOCAL_FILE,
        related_id="1",
        storage_key="",
    )
    llm_node.graph_runtime_state.variable_pool.add(["sys", "files"], file)

    result = llm_node._fetch_files(selector=["sys", "files"])
    assert result == [file]


def test_fetch_files_with_array_file_segment(llm_node):
    files = [
        File(
            id="1",
            tenant_id="test",
            type=FileType.IMAGE,
            filename="test1.jpg",
            transfer_method=FileTransferMethod.LOCAL_FILE,
            related_id="1",
            storage_key="",
        ),
        File(
            id="2",
            tenant_id="test",
            type=FileType.IMAGE,
            filename="test2.jpg",
            transfer_method=FileTransferMethod.LOCAL_FILE,
            related_id="2",
            storage_key="",
        ),
    ]
    llm_node.graph_runtime_state.variable_pool.add(["sys", "files"], ArrayFileSegment(value=files))

    result = llm_node._fetch_files(selector=["sys", "files"])
    assert result == files


def test_fetch_files_with_none_segment(llm_node):
    llm_node.graph_runtime_state.variable_pool.add(["sys", "files"], NoneSegment())

    result = llm_node._fetch_files(selector=["sys", "files"])
    assert result == []


def test_fetch_files_with_array_any_segment(llm_node):
    llm_node.graph_runtime_state.variable_pool.add(["sys", "files"], ArrayAnySegment(value=[]))

    result = llm_node._fetch_files(selector=["sys", "files"])
    assert result == []


def test_fetch_files_with_non_existent_variable(llm_node):
    result = llm_node._fetch_files(selector=["sys", "files"])
    assert result == []


# def test_fetch_prompt_messages__vison_disabled(faker, llm_node, model_config):
# TODO: Add test
# pass
# prompt_template = []
# llm_node.node_data.prompt_template = prompt_template

# fake_vision_detail = faker.random_element(
#     [ImagePromptMessageContent.DETAIL.HIGH, ImagePromptMessageContent.DETAIL.LOW]
# )
# fake_remote_url = faker.url()
# files = [
#     File(
#         id="1",
#         tenant_id="test",
#         type=FileType.IMAGE,
#         filename="test1.jpg",
#         transfer_method=FileTransferMethod.REMOTE_URL,
#         remote_url=fake_remote_url,
#         storage_key="",
#     )
# ]

# fake_query = faker.sentence()

# prompt_messages, _ = llm_node._fetch_prompt_messages(
#     sys_query=fake_query,
#     sys_files=files,
#     context=None,
#     memory=None,
#     model_config=model_config,
#     prompt_template=prompt_template,
#     memory_config=None,
#     vision_enabled=False,
#     vision_detail=fake_vision_detail,
#     variable_pool=llm_node.graph_runtime_state.variable_pool,
#     jinja2_variables=[],
# )

# assert prompt_messages == [UserPromptMessage(content=fake_query)]


# def test_fetch_prompt_messages__basic(faker, llm_node, model_config):
# TODO: Add test
# pass
# Setup dify config
# dify_config.MULTIMODAL_SEND_FORMAT = "url"

# # Generate fake values for prompt template
# fake_assistant_prompt = faker.sentence()
# fake_query = faker.sentence()
# fake_context = faker.sentence()
# fake_window_size = faker.random_int(min=1, max=3)
# fake_vision_detail = faker.random_element(
#     [ImagePromptMessageContent.DETAIL.HIGH, ImagePromptMessageContent.DETAIL.LOW]
# )
# fake_remote_url = faker.url()

# # Setup mock memory with history messages
# mock_history = [
#     UserPromptMessage(content=faker.sentence()),
#     AssistantPromptMessage(content=faker.sentence()),
#     UserPromptMessage(content=faker.sentence()),
#     AssistantPromptMessage(content=faker.sentence()),
#     UserPromptMessage(content=faker.sentence()),
#     AssistantPromptMessage(content=faker.sentence()),
# ]

# # Setup memory configuration
# memory_config = MemoryConfig(
#     role_prefix=MemoryConfig.RolePrefix(user="Human", assistant="Assistant"),
#     window=MemoryConfig.WindowConfig(enabled=True, size=fake_window_size),
#     query_prompt_template=None,
# )

# memory = MockTokenBufferMemory(history_messages=mock_history)

# # Test scenarios covering different file input combinations
# test_scenarios = [
#     LLMNodeTestScenario(
#         description="No files",
#         sys_query=fake_query,
#         sys_files=[],
#         features=[],
#         vision_enabled=False,
#         vision_detail=None,
#         window_size=fake_window_size,
#         prompt_template=[
#             LLMNodeChatModelMessage(
#                 text=fake_context,
#                 role=PromptMessageRole.SYSTEM,
#                 edition_type="basic",
#             ),
#             LLMNodeChatModelMessage(
#                 text="{#context#}",
#                 role=PromptMessageRole.USER,
#                 edition_type="basic",
#             ),
#             LLMNodeChatModelMessage(
#                 text=fake_assistant_prompt,
#                 role=PromptMessageRole.ASSISTANT,
#                 edition_type="basic",
#             ),
#         ],
#         expected_messages=[
#             SystemPromptMessage(content=fake_context),
#             UserPromptMessage(content=fake_context),
#             AssistantPromptMessage(content=fake_assistant_prompt),
#         ]
#         + mock_history[fake_window_size * -2 :]
#         + [
#             UserPromptMessage(content=fake_query),
#         ],
#     ),
#     LLMNodeTestScenario(
#         description="User files",
#         sys_query=fake_query,
#         sys_files=[
#             File(
#                 tenant_id="test",
#                 type=FileType.IMAGE,
#                 filename="test1.jpg",
#                 transfer_method=FileTransferMethod.REMOTE_URL,
#                 remote_url=fake_remote_url,
#                 extension=".jpg",
#                 mime_type="image/jpg",
#                 storage_key="",
#             )
#         ],
#         vision_enabled=True,
#         vision_detail=fake_vision_detail,
#         features=[ModelFeature.VISION],
#         window_size=fake_window_size,
#         prompt_template=[
#             LLMNodeChatModelMessage(
#                 text=fake_context,
#                 role=PromptMessageRole.SYSTEM,
#                 edition_type="basic",
#             ),
#             LLMNodeChatModelMessage(
#                 text="{#context#}",
#                 role=PromptMessageRole.USER,
#                 edition_type="basic",
#             ),
#             LLMNodeChatModelMessage(
#                 text=fake_assistant_prompt,
#                 role=PromptMessageRole.ASSISTANT,
#                 edition_type="basic",
#             ),
#         ],
#         expected_messages=[
#             SystemPromptMessage(content=fake_context),
#             UserPromptMessage(content=fake_context),
#             AssistantPromptMessage(content=fake_assistant_prompt),
#         ]
#         + mock_history[fake_window_size * -2 :]
#         + [
#             UserPromptMessage(
#                 content=[
#                     TextPromptMessageContent(data=fake_query),
#                     ImagePromptMessageContent(
#                         url=fake_remote_url, mime_type="image/jpg", format="jpg", detail=fake_vision_detail
#                     ),
#                 ]
#             ),
#         ],
#     ),
#     LLMNodeTestScenario(
#         description="Prompt template with variable selector of File",
#         sys_query=fake_query,
#         sys_files=[],
#         vision_enabled=False,
#         vision_detail=fake_vision_detail,
#         features=[ModelFeature.VISION],
#         window_size=fake_window_size,
#         prompt_template=[
#             LLMNodeChatModelMessage(
#                 text="{{#input.image#}}",
#                 role=PromptMessageRole.USER,
#                 edition_type="basic",
#             ),
#         ],
#         expected_messages=[
#             UserPromptMessage(
#                 content=[
#                     ImagePromptMessageContent(
#                         url=fake_remote_url, mime_type="image/jpg", format="jpg", detail=fake_vision_detail
#                     ),
#                 ]
#             ),
#         ]
#         + mock_history[fake_window_size * -2 :]
#         + [UserPromptMessage(content=fake_query)],
#         file_variables={
#             "input.image": File(
#                 tenant_id="test",
#                 type=FileType.IMAGE,
#                 filename="test1.jpg",
#                 transfer_method=FileTransferMethod.REMOTE_URL,
#                 remote_url=fake_remote_url,
#                 extension=".jpg",
#                 mime_type="image/jpg",
#                 storage_key="",
#             )
#         },
#     ),
# ]

# for scenario in test_scenarios:
#     model_config.model_schema.features = scenario.features

#     for k, v in scenario.file_variables.items():
#         selector = k.split(".")
#         llm_node.graph_runtime_state.variable_pool.add(selector, v)

#     # Call the method under test
#     prompt_messages, _ = llm_node._fetch_prompt_messages(
#         sys_query=scenario.sys_query,
#         sys_files=scenario.sys_files,
#         context=fake_context,
#         memory=memory,
#         model_config=model_config,
#         prompt_template=scenario.prompt_template,
#         memory_config=memory_config,
#         vision_enabled=scenario.vision_enabled,
#         vision_detail=scenario.vision_detail,
#         variable_pool=llm_node.graph_runtime_state.variable_pool,
#         jinja2_variables=[],
#     )

# # Verify the result
# assert len(prompt_messages) == len(scenario.expected_messages), f"Scenario failed: {scenario.description}"
# assert prompt_messages == scenario.expected_messages, (
#     f"Message content mismatch in scenario: {scenario.description}"
# )


def test_handle_list_messages_basic(llm_node):
    messages = [
        LLMNodeChatModelMessage(
            text="Hello, {#context#}",
            role=PromptMessageRole.USER,
            edition_type="basic",
        )
    ]
    context = "world"
    jinja2_variables = []
    variable_pool = llm_node.graph_runtime_state.variable_pool
    vision_detail_config = ImagePromptMessageContent.DETAIL.HIGH

    result = llm_node._handle_list_messages(
        messages=messages,
        context=context,
        jinja2_variables=jinja2_variables,
        variable_pool=variable_pool,
        vision_detail_config=vision_detail_config,
    )

    assert len(result) == 1
    assert isinstance(result[0], UserPromptMessage)
    assert result[0].content == [TextPromptMessageContent(data="Hello, world")]
