import dataclasses
import secrets
import uuid
from unittest.mock import MagicMock, Mock, patch

import pytest
from sqlalchemy import Engine
from sqlalchemy.orm import Session

from dify_graph.constants import SYSTEM_VARIABLE_NODE_ID
from dify_graph.enums import BuiltinNodeTypes, SystemVariableKey
from dify_graph.variables.segments import StringSegment
from dify_graph.variables.types import SegmentType
from libs.uuid_utils import uuidv7
from models.account import Account
from models.enums import DraftVariableType
from models.workflow import (
    Workflow,
    WorkflowDraftVariable,
    WorkflowDraftVariableFile,
    WorkflowNodeExecutionModel,
    is_system_variable_editable,
)
from services.workflow_draft_variable_service import (
    DraftVariableSaver,
    VariableResetError,
    WorkflowDraftVariableService,
)


@pytest.fixture
def mock_engine() -> Engine:
    return Mock(spec=Engine)


@pytest.fixture
def mock_session(mock_engine) -> Session:
    mock_session = Mock(spec=Session)
    mock_session.get_bind.return_value = mock_engine
    return mock_session


class TestDraftVariableSaver:
    def _get_test_app_id(self):
        suffix = secrets.token_hex(6)
        return f"test_app_id_{suffix}"

    def test__should_variable_be_visible(self):
        mock_session = MagicMock(spec=Session)
        mock_user = Account(name="test", email="test@example.com")
        mock_user.id = str(uuid.uuid4())
        test_app_id = self._get_test_app_id()
        saver = DraftVariableSaver(
            session=mock_session,
            app_id=test_app_id,
            node_id="test_node_id",
            node_type=BuiltinNodeTypes.START,
            node_execution_id="test_execution_id",
            user=mock_user,
        )
        assert saver._should_variable_be_visible("123_456", BuiltinNodeTypes.IF_ELSE, "output") == False
        assert saver._should_variable_be_visible("123", BuiltinNodeTypes.START, "output") == True

    def test__normalize_variable_for_start_node(self):
        @dataclasses.dataclass(frozen=True)
        class TestCase:
            name: str
            input_node_id: str
            input_name: str
            expected_node_id: str
            expected_name: str

        _NODE_ID = "1747228642872"
        cases = [
            TestCase(
                name="name with `sys.` prefix should return the system node_id",
                input_node_id=_NODE_ID,
                input_name="sys.workflow_id",
                expected_node_id=SYSTEM_VARIABLE_NODE_ID,
                expected_name="workflow_id",
            ),
            TestCase(
                name="name without `sys.` prefix should return the original input node_id",
                input_node_id=_NODE_ID,
                input_name="start_input",
                expected_node_id=_NODE_ID,
                expected_name="start_input",
            ),
            TestCase(
                name="dummy_variable should return the original input node_id",
                input_node_id=_NODE_ID,
                input_name="__dummy__",
                expected_node_id=_NODE_ID,
                expected_name="__dummy__",
            ),
        ]

        mock_session = MagicMock(spec=Session)
        mock_user = MagicMock()
        test_app_id = self._get_test_app_id()
        saver = DraftVariableSaver(
            session=mock_session,
            app_id=test_app_id,
            node_id=_NODE_ID,
            node_type=BuiltinNodeTypes.START,
            node_execution_id="test_execution_id",
            user=mock_user,
        )
        for idx, c in enumerate(cases, 1):
            fail_msg = f"Test case {c.name} failed, index={idx}"
            node_id, name = saver._normalize_variable_for_start_node(c.input_name)
            assert node_id == c.expected_node_id, fail_msg
            assert name == c.expected_name, fail_msg

    @pytest.fixture
    def mock_session(self):
        """Mock SQLAlchemy session."""
        from sqlalchemy import Engine

        mock_session = MagicMock(spec=Session)
        mock_engine = MagicMock(spec=Engine)
        mock_session.get_bind.return_value = mock_engine
        return mock_session

    @pytest.fixture
    def draft_saver(self, mock_session):
        """Create DraftVariableSaver instance with user context."""
        # Create a mock user
        mock_user = MagicMock(spec=Account)
        mock_user.id = "test-user-id"
        mock_user.tenant_id = "test-tenant-id"

        return DraftVariableSaver(
            session=mock_session,
            app_id="test-app-id",
            node_id="test-node-id",
            node_type=BuiltinNodeTypes.LLM,
            node_execution_id="test-execution-id",
            user=mock_user,
        )

    def test_draft_saver_with_small_variables(self, draft_saver, mock_session):
        with patch(
            "services.workflow_draft_variable_service.DraftVariableSaver._try_offload_large_variable", autospec=True
        ) as _mock_try_offload:
            _mock_try_offload.return_value = None
            mock_segment = StringSegment(value="small value")
            draft_var = draft_saver._create_draft_variable(name="small_var", value=mock_segment, visible=True)

            # Should not have large variable metadata
            assert draft_var.file_id is None
            _mock_try_offload.return_value = None

    def test_draft_saver_with_large_variables(self, draft_saver, mock_session):
        with patch(
            "services.workflow_draft_variable_service.DraftVariableSaver._try_offload_large_variable", autospec=True
        ) as _mock_try_offload:
            mock_segment = StringSegment(value="small value")
            mock_draft_var_file = WorkflowDraftVariableFile(
                id=str(uuidv7()),
                size=1024,
                length=10,
                value_type=SegmentType.ARRAY_STRING,
                upload_file_id=str(uuid.uuid4()),
            )

            _mock_try_offload.return_value = mock_segment, mock_draft_var_file
            draft_var = draft_saver._create_draft_variable(name="small_var", value=mock_segment, visible=True)

            # Should not have large variable metadata
            assert draft_var.file_id == mock_draft_var_file.id

    @patch("services.workflow_draft_variable_service._batch_upsert_draft_variable", autospec=True)
    def test_save_method_integration(self, mock_batch_upsert, draft_saver):
        """Test complete save workflow."""
        outputs = {"result": {"data": "test_output"}, "metadata": {"type": "llm_response"}}

        draft_saver.save(outputs=outputs)

        # Should batch upsert draft variables
        mock_batch_upsert.assert_called_once()
        draft_vars = mock_batch_upsert.call_args[0][1]
        assert len(draft_vars) == 2

    @patch("services.workflow_draft_variable_service._batch_upsert_draft_variable", autospec=True)
    def test_start_node_save_persists_sys_timestamp_and_workflow_run_id(self, mock_batch_upsert):
        """Start node should persist common `sys.*` variables, not only `sys.files`."""
        mock_session = MagicMock(spec=Session)
        mock_user = MagicMock(spec=Account)
        mock_user.id = "test-user-id"
        mock_user.tenant_id = "test-tenant-id"

        saver = DraftVariableSaver(
            session=mock_session,
            app_id="test-app-id",
            node_id="start-node-id",
            node_type=BuiltinNodeTypes.START,
            node_execution_id="exec-id",
            user=mock_user,
        )

        outputs = {
            f"{SYSTEM_VARIABLE_NODE_ID}.{SystemVariableKey.TIMESTAMP}": 1700000000,
            f"{SYSTEM_VARIABLE_NODE_ID}.{SystemVariableKey.WORKFLOW_EXECUTION_ID}": "run-id-123",
        }

        saver.save(outputs=outputs)

        mock_batch_upsert.assert_called_once()
        draft_vars = mock_batch_upsert.call_args[0][1]

        # plus one dummy output because there are no non-sys Start inputs
        assert len(draft_vars) == 3

        sys_vars = [v for v in draft_vars if v.node_id == SYSTEM_VARIABLE_NODE_ID]
        assert {v.name for v in sys_vars} == {
            str(SystemVariableKey.TIMESTAMP),
            str(SystemVariableKey.WORKFLOW_EXECUTION_ID),
        }


class TestWorkflowDraftVariableService:
    def _get_test_app_id(self):
        suffix = secrets.token_hex(6)
        return f"test_app_id_{suffix}"

    def _create_test_workflow(self, app_id: str) -> Workflow:
        """Create a real Workflow instance for testing"""
        return Workflow.new(
            tenant_id="test_tenant_id",
            app_id=app_id,
            type="workflow",
            version="draft",
            graph='{"nodes": [], "edges": []}',
            features="{}",
            created_by="test_user_id",
            environment_variables=[],
            conversation_variables=[],
            rag_pipeline_variables=[],
        )

    def test_reset_conversation_variable(self, mock_session):
        """Test resetting a conversation variable"""
        service = WorkflowDraftVariableService(mock_session)

        test_app_id = self._get_test_app_id()
        workflow = self._create_test_workflow(test_app_id)

        # Create real conversation variable
        test_value = StringSegment(value="test_value")
        variable = WorkflowDraftVariable.new_conversation_variable(
            app_id=test_app_id, name="test_var", value=test_value, description="Test conversation variable"
        )

        # Mock the _reset_conv_var method
        expected_result = WorkflowDraftVariable.new_conversation_variable(
            app_id=test_app_id,
            name="test_var",
            value=StringSegment(value="reset_value"),
        )
        with patch.object(service, "_reset_conv_var", return_value=expected_result, autospec=True) as mock_reset_conv:
            result = service.reset_variable(workflow, variable)

            mock_reset_conv.assert_called_once_with(workflow, variable)
            assert result == expected_result

    def test_reset_node_variable_with_no_execution_id(self, mock_session):
        """Test resetting a node variable with no execution ID - should delete variable"""
        service = WorkflowDraftVariableService(mock_session)

        test_app_id = self._get_test_app_id()
        workflow = self._create_test_workflow(test_app_id)

        # Create real node variable with no execution ID
        test_value = StringSegment(value="test_value")
        variable = WorkflowDraftVariable.new_node_variable(
            app_id=test_app_id,
            node_id="test_node_id",
            name="test_var",
            value=test_value,
            node_execution_id="exec-id",  # Set initially
        )
        # Manually set to None to simulate the test condition
        variable.node_execution_id = None

        result = service._reset_node_var_or_sys_var(workflow, variable)

        # Should delete the variable and return None
        mock_session.delete.assert_called_once_with(instance=variable)
        mock_session.flush.assert_called_once()
        assert result is None

    def test_reset_node_variable_with_missing_execution_record(
        self,
        mock_engine,
        mock_session,
        monkeypatch,
    ):
        """Test resetting a node variable when execution record doesn't exist"""
        mock_repo_session = Mock(spec=Session)

        mock_session_maker = MagicMock()
        # Mock the context manager protocol for sessionmaker
        mock_session_maker.return_value.__enter__.return_value = mock_repo_session
        mock_session_maker.return_value.__exit__.return_value = None
        monkeypatch.setattr("services.workflow_draft_variable_service.sessionmaker", mock_session_maker)
        service = WorkflowDraftVariableService(mock_session)

        # Mock the repository to return None (no execution record found)
        service._api_node_execution_repo = Mock()
        service._api_node_execution_repo.get_execution_by_id.return_value = None

        test_app_id = self._get_test_app_id()
        workflow = self._create_test_workflow(test_app_id)

        # Create real node variable with execution ID
        test_value = StringSegment(value="test_value")
        variable = WorkflowDraftVariable.new_node_variable(
            app_id=test_app_id, node_id="test_node_id", name="test_var", value=test_value, node_execution_id="exec-id"
        )
        # Variable is editable by default from factory method

        result = service._reset_node_var_or_sys_var(workflow, variable)

        mock_session_maker.assert_called_once_with(bind=mock_engine, expire_on_commit=False)
        # Should delete the variable and return None
        mock_session.delete.assert_called_once_with(instance=variable)
        mock_session.flush.assert_called_once()
        assert result is None

    def test_reset_node_variable_with_valid_execution_record(
        self,
        mock_session,
        monkeypatch,
    ):
        """Test resetting a node variable with valid execution record - should restore from execution"""
        mock_repo_session = Mock(spec=Session)

        mock_session_maker = MagicMock()
        # Mock the context manager protocol for sessionmaker
        mock_session_maker.return_value.__enter__.return_value = mock_repo_session
        mock_session_maker.return_value.__exit__.return_value = None
        mock_session_maker = monkeypatch.setattr(
            "services.workflow_draft_variable_service.sessionmaker", mock_session_maker
        )
        service = WorkflowDraftVariableService(mock_session)

        # Create mock execution record
        mock_execution = Mock(spec=WorkflowNodeExecutionModel)
        mock_execution.load_full_outputs.return_value = {"test_var": "output_value"}

        # Mock the repository to return the execution record
        service._api_node_execution_repo = Mock()
        service._api_node_execution_repo.get_execution_by_id.return_value = mock_execution

        test_app_id = self._get_test_app_id()
        workflow = self._create_test_workflow(test_app_id)

        # Create real node variable with execution ID
        test_value = StringSegment(value="original_value")
        variable = WorkflowDraftVariable.new_node_variable(
            app_id=test_app_id, node_id="test_node_id", name="test_var", value=test_value, node_execution_id="exec-id"
        )
        # Variable is editable by default from factory method

        # Mock workflow methods
        mock_node_config = {"type": "test_node"}
        with (
            patch.object(workflow, "get_node_config_by_id", return_value=mock_node_config, autospec=True),
            patch.object(workflow, "get_node_type_from_node_config", return_value=BuiltinNodeTypes.LLM, autospec=True),
        ):
            result = service._reset_node_var_or_sys_var(workflow, variable)

            # Verify last_edited_at was reset
            assert variable.last_edited_at is None
            # Verify session.flush was called
            mock_session.flush.assert_called()

            # Should return the updated variable
            assert result == variable

    def test_reset_non_editable_system_variable_raises_error(self, mock_session):
        """Test that resetting a non-editable system variable raises an error"""
        service = WorkflowDraftVariableService(mock_session)

        test_app_id = self._get_test_app_id()
        workflow = self._create_test_workflow(test_app_id)

        # Create a non-editable system variable (workflow_id is not editable)
        test_value = StringSegment(value="test_workflow_id")
        variable = WorkflowDraftVariable.new_sys_variable(
            app_id=test_app_id,
            name="workflow_id",  # This is not in _EDITABLE_SYSTEM_VARIABLE
            value=test_value,
            node_execution_id="exec-id",
            editable=False,  # Non-editable system variable
        )

        with pytest.raises(VariableResetError) as exc_info:
            service.reset_variable(workflow, variable)
        assert "cannot reset system variable" in str(exc_info.value)
        assert f"variable_id={variable.id}" in str(exc_info.value)

    def test_reset_editable_system_variable_succeeds(self, mock_session):
        """Test that resetting an editable system variable succeeds"""
        service = WorkflowDraftVariableService(mock_session)

        test_app_id = self._get_test_app_id()
        workflow = self._create_test_workflow(test_app_id)

        # Create an editable system variable (files is editable)
        test_value = StringSegment(value="[]")
        variable = WorkflowDraftVariable.new_sys_variable(
            app_id=test_app_id,
            name="files",  # This is in _EDITABLE_SYSTEM_VARIABLE
            value=test_value,
            node_execution_id="exec-id",
            editable=True,  # Editable system variable
        )

        # Create mock execution record
        mock_execution = Mock(spec=WorkflowNodeExecutionModel)
        mock_execution.load_full_outputs.return_value = {"sys.files": "[]"}

        # Mock the repository to return the execution record
        service._api_node_execution_repo = Mock()
        service._api_node_execution_repo.get_execution_by_id.return_value = mock_execution

        result = service._reset_node_var_or_sys_var(workflow, variable)

        # Should succeed and return the variable
        assert result == variable
        assert variable.last_edited_at is None
        mock_session.flush.assert_called()

    def test_reset_query_system_variable_succeeds(self, mock_session):
        """Test that resetting query system variable (another editable one) succeeds"""
        service = WorkflowDraftVariableService(mock_session)

        test_app_id = self._get_test_app_id()
        workflow = self._create_test_workflow(test_app_id)

        # Create an editable system variable (query is editable)
        test_value = StringSegment(value="original query")
        variable = WorkflowDraftVariable.new_sys_variable(
            app_id=test_app_id,
            name="query",  # This is in _EDITABLE_SYSTEM_VARIABLE
            value=test_value,
            node_execution_id="exec-id",
            editable=True,  # Editable system variable
        )

        # Create mock execution record
        mock_execution = Mock(spec=WorkflowNodeExecutionModel)
        mock_execution.load_full_outputs.return_value = {"sys.query": "reset query"}

        # Mock the repository to return the execution record
        service._api_node_execution_repo = Mock()
        service._api_node_execution_repo.get_execution_by_id.return_value = mock_execution

        result = service._reset_node_var_or_sys_var(workflow, variable)

        # Should succeed and return the variable
        assert result == variable
        assert variable.last_edited_at is None
        mock_session.flush.assert_called()

    def test_system_variable_editability_check(self):
        """Test the system variable editability function directly"""
        # Test editable system variables
        assert is_system_variable_editable("files") == True
        assert is_system_variable_editable("query") == True

        # Test non-editable system variables
        assert is_system_variable_editable("workflow_id") == False
        assert is_system_variable_editable("conversation_id") == False
        assert is_system_variable_editable("user_id") == False

    def test_workflow_draft_variable_factory_methods(self):
        """Test that factory methods create proper instances"""
        test_app_id = self._get_test_app_id()
        test_value = StringSegment(value="test_value")

        # Test conversation variable factory
        conv_var = WorkflowDraftVariable.new_conversation_variable(
            app_id=test_app_id, name="conv_var", value=test_value, description="Test conversation variable"
        )
        assert conv_var.get_variable_type() == DraftVariableType.CONVERSATION
        assert conv_var.editable == True
        assert conv_var.node_execution_id is None

        # Test system variable factory
        sys_var = WorkflowDraftVariable.new_sys_variable(
            app_id=test_app_id, name="workflow_id", value=test_value, node_execution_id="exec-id", editable=False
        )
        assert sys_var.get_variable_type() == DraftVariableType.SYS
        assert sys_var.editable == False
        assert sys_var.node_execution_id == "exec-id"

        # Test node variable factory
        node_var = WorkflowDraftVariable.new_node_variable(
            app_id=test_app_id,
            node_id="node-id",
            name="node_var",
            value=test_value,
            node_execution_id="exec-id",
            visible=True,
            editable=True,
        )
        assert node_var.get_variable_type() == DraftVariableType.NODE
        assert node_var.visible == True
        assert node_var.editable == True
        assert node_var.node_execution_id == "exec-id"
