import os
import time
import uuid
from unittest.mock import MagicMock

import pytest

from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom
from core.model_manager import ModelInstance
from core.workflow.node_runtime import DifyPromptMessageSerializer
from core.workflow.system_variables import build_system_variables
from extensions.ext_database import db
from graphon.enums import WorkflowNodeExecutionStatus
from graphon.model_runtime.entities import AssistantPromptMessage, UserPromptMessage
from graphon.nodes.llm.protocols import CredentialsProvider, ModelFactory
from graphon.nodes.parameter_extractor.entities import ParameterExtractorNodeData
from graphon.nodes.parameter_extractor.parameter_extractor_node import ParameterExtractorNode
from graphon.runtime import GraphRuntimeState, VariablePool
from tests.integration_tests.workflow.nodes.__mock.model import get_mocked_fetch_model_instance
from tests.workflow_test_utils import build_test_graph_init_params

pytest_plugins = ("tests.integration_tests.model_runtime.__mock.plugin_daemon",)


def get_mocked_fetch_memory(memory_text: str):
    class MemoryMock:
        def get_history_prompt_messages(
            self,
            max_token_limit: int = 2000,
            message_limit: int | None = None,
        ):
            return [UserPromptMessage(content=memory_text), AssistantPromptMessage(content="mocked answer")]

    return MagicMock(return_value=MemoryMock())


def init_parameter_extractor_node(config: dict, memory=None):
    graph_config = {
        "edges": [
            {
                "id": "start-source-next-target",
                "source": "start",
                "target": "llm",
            },
        ],
        "nodes": [{"data": {"type": "start", "title": "Start"}, "id": "start"}, config],
    }

    init_params = build_test_graph_init_params(
        workflow_id="1",
        graph_config=graph_config,
        tenant_id="1",
        app_id="1",
        user_id="1",
        user_from=UserFrom.ACCOUNT,
        invoke_from=InvokeFrom.DEBUGGER,
        call_depth=0,
    )

    # construct variable pool
    variable_pool = VariablePool.from_bootstrap(
        system_variables=build_system_variables(
            user_id="aaa", files=[], query="what's the weather in SF", conversation_id="abababa"
        ),
        user_inputs={},
        environment_variables=[],
        conversation_variables=[],
    )
    variable_pool.add(["a", "args1"], 1)
    variable_pool.add(["a", "args2"], 2)

    graph_runtime_state = GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter())

    node = ParameterExtractorNode(
        node_id=str(uuid.uuid4()),
        data=ParameterExtractorNodeData.model_validate(config["data"]),
        graph_init_params=init_params,
        graph_runtime_state=graph_runtime_state,
        credentials_provider=MagicMock(spec=CredentialsProvider),
        model_factory=MagicMock(spec=ModelFactory),
        model_instance=MagicMock(spec=ModelInstance),
        memory=memory,
        prompt_message_serializer=DifyPromptMessageSerializer(),
    )
    return node


def _mock_db_session_close(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(db.session, "close", MagicMock())


def test_function_calling_parameter_extractor(setup_model_mock, monkeypatch: pytest.MonkeyPatch):
    """
    Test function calling for parameter extractor.
    """
    node = init_parameter_extractor_node(
        config={
            "id": "llm",
            "data": {
                "title": "123",
                "type": "parameter-extractor",
                "model": {
                    "provider": "langgenius/openai/openai",
                    "name": "gpt-3.5-turbo",
                    "mode": "chat",
                    "completion_params": {},
                },
                "query": ["sys", "query"],
                "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}],
                "instruction": "",
                "reasoning_mode": "function_call",
                "memory": None,
            },
        }
    )

    node._model_instance = get_mocked_fetch_model_instance(
        provider="langgenius/openai/openai",
        model="gpt-3.5-turbo",
        mode="chat",
        credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")},
    )()
    _mock_db_session_close(monkeypatch)

    result = node._run()

    assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED
    assert result.outputs is not None
    assert result.outputs.get("location") == "kawaii"
    assert result.outputs.get("__reason") == None


def test_instructions(setup_model_mock, monkeypatch: pytest.MonkeyPatch):
    """
    Test chat parameter extractor.
    """
    node = init_parameter_extractor_node(
        config={
            "id": "llm",
            "data": {
                "title": "123",
                "type": "parameter-extractor",
                "model": {
                    "provider": "langgenius/openai/openai",
                    "name": "gpt-3.5-turbo",
                    "mode": "chat",
                    "completion_params": {},
                },
                "query": ["sys", "query"],
                "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}],
                "reasoning_mode": "function_call",
                "instruction": "{{#sys.query#}}",
                "memory": None,
            },
        },
    )

    node._model_instance = get_mocked_fetch_model_instance(
        provider="langgenius/openai/openai",
        model="gpt-3.5-turbo",
        mode="chat",
        credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")},
    )()
    _mock_db_session_close(monkeypatch)

    result = node._run()

    assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED
    assert result.outputs is not None
    assert result.outputs.get("location") == "kawaii"
    assert result.outputs.get("__reason") == None

    process_data = result.process_data

    assert process_data is not None
    process_data.get("prompts")

    for prompt in process_data.get("prompts", []):
        if prompt.get("role") == "system":
            assert "what's the weather in SF" in prompt.get("text")


def test_chat_parameter_extractor(setup_model_mock, monkeypatch: pytest.MonkeyPatch):
    """
    Test chat parameter extractor.
    """
    node = init_parameter_extractor_node(
        config={
            "id": "llm",
            "data": {
                "title": "123",
                "type": "parameter-extractor",
                "model": {
                    "provider": "langgenius/openai/openai",
                    "name": "gpt-3.5-turbo",
                    "mode": "chat",
                    "completion_params": {},
                },
                "query": ["sys", "query"],
                "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}],
                "reasoning_mode": "prompt",
                "instruction": "",
                "memory": None,
            },
        },
    )

    node._model_instance = get_mocked_fetch_model_instance(
        provider="langgenius/openai/openai",
        model="gpt-3.5-turbo",
        mode="chat",
        credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")},
    )()
    _mock_db_session_close(monkeypatch)

    result = node._run()

    assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED
    assert result.outputs is not None
    assert result.outputs.get("location") == ""
    assert (
        result.outputs.get("__reason")
        == "Failed to extract result from function call or text response, using empty result."
    )
    assert result.process_data is not None
    prompts = result.process_data.get("prompts", [])

    for prompt in prompts:
        if prompt.get("role") == "user":
            if "<structure>" in prompt.get("text"):
                assert '<structure>\n{"type": "object"' in prompt.get("text")


def test_completion_parameter_extractor(setup_model_mock, monkeypatch: pytest.MonkeyPatch):
    """
    Test completion parameter extractor.
    """
    node = init_parameter_extractor_node(
        config={
            "id": "llm",
            "data": {
                "title": "123",
                "type": "parameter-extractor",
                "model": {
                    "provider": "langgenius/openai/openai",
                    "name": "gpt-3.5-turbo-instruct",
                    "mode": "completion",
                    "completion_params": {},
                },
                "query": ["sys", "query"],
                "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}],
                "reasoning_mode": "prompt",
                "instruction": "{{#sys.query#}}",
                "memory": None,
            },
        },
    )

    node._model_instance = get_mocked_fetch_model_instance(
        provider="langgenius/openai/openai",
        model="gpt-3.5-turbo-instruct",
        mode="completion",
        credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")},
    )()
    _mock_db_session_close(monkeypatch)

    result = node._run()

    assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED
    assert result.outputs is not None
    assert result.outputs.get("location") == ""
    assert (
        result.outputs.get("__reason")
        == "Failed to extract result from function call or text response, using empty result."
    )
    assert result.process_data is not None
    assert len(result.process_data.get("prompts", [])) == 1
    assert "SF" in result.process_data.get("prompts", [])[0].get("text")


def test_extract_json_response():
    """
    Test extract json response.
    """

    node = init_parameter_extractor_node(
        config={
            "id": "llm",
            "data": {
                "title": "123",
                "type": "parameter-extractor",
                "model": {
                    "provider": "langgenius/openai/openai",
                    "name": "gpt-3.5-turbo-instruct",
                    "mode": "completion",
                    "completion_params": {},
                },
                "query": ["sys", "query"],
                "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}],
                "reasoning_mode": "prompt",
                "instruction": "{{#sys.query#}}",
                "memory": None,
            },
        },
    )

    result = node._extract_complete_json_response("""
        uwu{ovo}
        {
            "location": "kawaii"
        }
        hello world.
    """)

    assert result is not None
    assert result["location"] == "kawaii"


def test_extract_json_from_tool_call():
    """
    Test extract json response.
    """

    node = init_parameter_extractor_node(
        config={
            "id": "llm",
            "data": {
                "title": "123",
                "type": "parameter-extractor",
                "model": {
                    "provider": "langgenius/openai/openai",
                    "name": "gpt-3.5-turbo-instruct",
                    "mode": "completion",
                    "completion_params": {},
                },
                "query": ["sys", "query"],
                "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}],
                "reasoning_mode": "prompt",
                "instruction": "{{#sys.query#}}",
                "memory": None,
            },
        },
    )

    result = node._extract_json_from_tool_call(
        AssistantPromptMessage.ToolCall(
            id="llm",
            type="parameter-extractor",
            function=AssistantPromptMessage.ToolCall.ToolCallFunction(
                name="foo", arguments="""{"location":"kawaii"}{"location": 1}"""
            ),
        )
    )

    assert result is not None
    assert result["location"] == "kawaii"


def test_chat_parameter_extractor_with_memory(setup_model_mock, monkeypatch: pytest.MonkeyPatch):
    """
    Test chat parameter extractor with memory.
    """
    node = init_parameter_extractor_node(
        config={
            "id": "llm",
            "data": {
                "title": "123",
                "type": "parameter-extractor",
                "model": {
                    "provider": "langgenius/openai/openai",
                    "name": "gpt-3.5-turbo",
                    "mode": "chat",
                    "completion_params": {},
                },
                "query": ["sys", "query"],
                "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}],
                "reasoning_mode": "prompt",
                "instruction": "",
                "memory": {"window": {"enabled": True, "size": 50}},
            },
        },
        memory=get_mocked_fetch_memory("customized memory")(),
    )

    node._model_instance = get_mocked_fetch_model_instance(
        provider="langgenius/openai/openai",
        model="gpt-3.5-turbo",
        mode="chat",
        credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")},
    )()
    _mock_db_session_close(monkeypatch)

    result = node._run()

    assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED
    assert result.outputs is not None
    assert result.outputs.get("location") == ""
    assert (
        result.outputs.get("__reason")
        == "Failed to extract result from function call or text response, using empty result."
    )
    assert result.process_data is not None
    prompts = result.process_data.get("prompts", [])

    latest_role = None
    for prompt in prompts:
        if prompt.get("role") == "user":
            if "<structure>" in prompt.get("text"):
                assert '<structure>\n{"type": "object"' in prompt.get("text")
        elif prompt.get("role") == "system":
            assert "customized memory" in prompt.get("text")

        if latest_role is not None:
            assert latest_role != prompt.get("role")

        if prompt.get("role") in {"user", "assistant"}:
            latest_role = prompt.get("role")
