import json
import uuid
from collections import defaultdict

import pytest

from core.workflow.system_variables import build_system_variables, system_variables_to_mapping
from core.workflow.variable_pool_initializer import add_variables_to_pool
from core.workflow.variable_prefixes import (
    CONVERSATION_VARIABLE_NODE_ID,
    ENVIRONMENT_VARIABLE_NODE_ID,
    SYSTEM_VARIABLE_NODE_ID,
)
from factories.variable_factory import build_segment, segment_to_variable
from graphon.file import File, FileTransferMethod, FileType
from graphon.runtime import VariablePool
from graphon.variables import FileSegment, StringSegment
from graphon.variables.segments import (
    ArrayAnySegment,
    ArrayFileSegment,
    ArrayNumberSegment,
    ArrayObjectSegment,
    ArrayStringSegment,
    FloatSegment,
    IntegerSegment,
    NoneSegment,
    ObjectSegment,
)
from graphon.variables.variables import (
    ArrayNumberVariable,
    ArrayObjectVariable,
    ArrayStringVariable,
    FloatVariable,
    IntegerVariable,
    ObjectVariable,
    StringVariable,
    Variable,
)
from models.utils.file_input_compat import rebuild_serialized_graph_files_without_lookup


@pytest.fixture
def pool():
    variable_pool = VariablePool()
    add_variables_to_pool(
        variable_pool,
        build_system_variables(
            user_id="test_user_id",
            app_id="test_app_id",
            workflow_id="test_workflow_id",
        ),
    )
    return variable_pool


@pytest.fixture
def file():
    return File(
        file_type=FileType.DOCUMENT,
        transfer_method=FileTransferMethod.LOCAL_FILE,
        related_id="test_related_id",
        remote_url="test_url",
        filename="test_file.txt",
        storage_key="",
    )


def test_get_file_attribute(pool, file):
    # Add a FileSegment to the pool
    pool.add(("node_1", "file_var"), FileSegment(value=file))

    # Test getting the 'name' attribute of the file
    result = pool.get(("node_1", "file_var", "name"))

    assert result is not None
    assert result.value == file.filename

    # Test getting a non-existent attribute
    result = pool.get(("node_1", "file_var", "non_existent_attr"))
    assert result is None


class TestVariablePool:
    def test_constructor(self):
        pool = VariablePool()
        assert pool.variable_dictionary == defaultdict(dict)

        complex_system_vars = build_system_variables(
            user_id="test_user_id", app_id="test_app_id", workflow_id="test_workflow_id"
        )
        add_variables_to_pool(pool, complex_system_vars)
        add_variables_to_pool(
            pool,
            [
                segment_to_variable(
                    segment=build_segment(1),
                    selector=[ENVIRONMENT_VARIABLE_NODE_ID, "env_var_1"],
                    name="env_var_1",
                ),
                segment_to_variable(
                    segment=build_segment("1"),
                    selector=[CONVERSATION_VARIABLE_NODE_ID, "conv_var_1"],
                    name="conv_var_1",
                ),
            ],
        )

        assert pool.get([SYSTEM_VARIABLE_NODE_ID, "user_id"]) is not None
        assert pool.get([ENVIRONMENT_VARIABLE_NODE_ID, "env_var_1"]) is not None
        assert pool.get([CONVERSATION_VARIABLE_NODE_ID, "conv_var_1"]) is not None

    def test_constructor_loads_legacy_bootstrap_kwargs(self):
        pool = VariablePool(
            system_variables=build_system_variables(user_id="test_user_id"),
            environment_variables=[StringVariable(name="env_var", value="env-value")],
            conversation_variables=[StringVariable(name="conv_var", value="conv-value")],
            user_inputs={"ignored": "value"},
        )

        system_value = pool.get([SYSTEM_VARIABLE_NODE_ID, "user_id"])
        environment_value = pool.get([ENVIRONMENT_VARIABLE_NODE_ID, "env_var"])
        conversation_value = pool.get([CONVERSATION_VARIABLE_NODE_ID, "conv_var"])

        assert system_value is not None
        assert system_value.value == "test_user_id"
        assert environment_value is not None
        assert environment_value.value == "env-value"
        assert conversation_value is not None
        assert conversation_value.value == "conv-value"
        assert "system_variables" not in pool.model_dump()

    def test_get_system_variables(self):
        sys_var = build_system_variables(
            user_id="test_user_id",
            app_id="test_app_id",
            workflow_id="test_workflow_id",
            workflow_execution_id="test_execution_123",
            query="test query",
            conversation_id="test_conv_id",
            dialogue_count=5,
        )
        pool = VariablePool()
        add_variables_to_pool(pool, sys_var)
        system_values = system_variables_to_mapping(sys_var)

        kv = [
            ("user_id", system_values["user_id"]),
            ("app_id", system_values["app_id"]),
            ("workflow_id", system_values["workflow_id"]),
            ("workflow_run_id", system_values["workflow_run_id"]),
            ("query", system_values["query"]),
            ("conversation_id", system_values["conversation_id"]),
            ("dialogue_count", system_values["dialogue_count"]),
        ]
        for key, expected_value in kv:
            segment = pool.get([SYSTEM_VARIABLE_NODE_ID, key])
            assert segment is not None
            assert segment.value == expected_value


class TestVariablePoolSerialization:
    """Test cases for VariablePool serialization and deserialization using Pydantic's built-in methods.

    These tests focus exclusively on serialization/deserialization logic to ensure that
    VariablePool data can be properly serialized to dictionaries/JSON and reconstructed
    while preserving all data integrity.
    """

    _NODE1_ID = "node_1"
    _NODE2_ID = "node_2"
    _NODE3_ID = "node_3"

    def _create_pool_without_file(self):
        # Create comprehensive system variables
        system_vars = build_system_variables(
            user_id="test_user_id",
            app_id="test_app_id",
            workflow_id="test_workflow_id",
            workflow_execution_id="test_execution_123",
            query="test query",
            conversation_id="test_conv_id",
            dialogue_count=5,
        )

        # Create environment variables with all types including ArrayFileVariable
        env_vars: list[Variable] = [
            StringVariable(
                id="env_string_id",
                name="env_string",
                value="env_string_value",
                selector=[ENVIRONMENT_VARIABLE_NODE_ID, "env_string"],
            ),
            IntegerVariable(
                id="env_integer_id",
                name="env_integer",
                value=1,
                selector=[ENVIRONMENT_VARIABLE_NODE_ID, "env_integer"],
            ),
            FloatVariable(
                id="env_float_id",
                name="env_float",
                value=1.0,
                selector=[ENVIRONMENT_VARIABLE_NODE_ID, "env_float"],
            ),
        ]

        # Create conversation variables with complex data
        conv_vars: list[Variable] = [
            StringVariable(
                id="conv_string_id",
                name="conv_string",
                value="conv_string_value",
                selector=[CONVERSATION_VARIABLE_NODE_ID, "conv_string"],
            ),
            IntegerVariable(
                id="conv_integer_id",
                name="conv_integer",
                value=1,
                selector=[CONVERSATION_VARIABLE_NODE_ID, "conv_integer"],
            ),
            FloatVariable(
                id="conv_float_id",
                name="conv_float",
                value=1.0,
                selector=[CONVERSATION_VARIABLE_NODE_ID, "conv_float"],
            ),
            ObjectVariable(
                id="conv_object_id",
                name="conv_object",
                value={"key": "value", "nested": {"data": 123}},
                selector=[CONVERSATION_VARIABLE_NODE_ID, "conv_object"],
            ),
            ArrayStringVariable(
                id="conv_array_string_id",
                name="conv_array_string",
                value=["conv_array_string_value"],
                selector=[CONVERSATION_VARIABLE_NODE_ID, "conv_array_string"],
            ),
            ArrayNumberVariable(
                id="conv_array_number_id",
                name="conv_array_number",
                value=[1, 1.0],
                selector=[CONVERSATION_VARIABLE_NODE_ID, "conv_array_number"],
            ),
            ArrayObjectVariable(
                id="conv_array_object_id",
                name="conv_array_object",
                value=[{"a": 1}, {"b": "2"}],
                selector=[CONVERSATION_VARIABLE_NODE_ID, "conv_array_object"],
            ),
        ]

        # Create comprehensive user inputs
        user_inputs = {
            "string_input": "test_value",
            "number_input": 42,
            "object_input": {"nested": {"key": "value"}},
            "array_input": ["item1", "item2", "item3"],
        }

        # Create VariablePool
        pool = VariablePool()
        add_variables_to_pool(pool, system_vars)
        add_variables_to_pool(pool, env_vars)
        add_variables_to_pool(pool, conv_vars)
        return pool

    def _add_node_data_to_pool(self, pool: VariablePool, with_file=False):
        test_file = File(
            file_type=FileType.DOCUMENT,
            transfer_method=FileTransferMethod.LOCAL_FILE,
            related_id="test_related_id",
            remote_url="test_url",
            filename="test_file.txt",
            storage_key="test_storage_key",
        )

        # Add various segment types to variable dictionary
        pool.add((self._NODE1_ID, "string_var"), StringSegment(value="test_string"))
        pool.add((self._NODE1_ID, "int_var"), IntegerSegment(value=123))
        pool.add((self._NODE1_ID, "float_var"), FloatSegment(value=45.67))
        pool.add((self._NODE1_ID, "object_var"), ObjectSegment(value={"test": "data"}))
        if with_file:
            pool.add((self._NODE1_ID, "file_var"), FileSegment(value=test_file))
        pool.add((self._NODE1_ID, "none_var"), NoneSegment())

        # Add array segments including ArrayFileVariable
        pool.add((self._NODE2_ID, "array_string"), ArrayStringSegment(value=["a", "b", "c"]))
        pool.add((self._NODE2_ID, "array_number"), ArrayNumberSegment(value=[1, 2, 3]))
        pool.add((self._NODE2_ID, "array_object"), ArrayObjectSegment(value=[{"a": 1}, {"b": 2}]))
        if with_file:
            pool.add((self._NODE2_ID, "array_file"), ArrayFileSegment(value=[test_file]))
        pool.add((self._NODE2_ID, "array_any"), ArrayAnySegment(value=["mixed", 123, {"key": "value"}]))

    def test_system_variables(self):
        sys_vars = build_system_variables(
            user_id="test_user_id",
            app_id="test_app_id",
            workflow_id="test_workflow_id",
            workflow_execution_id="test_execution_123",
            query="test query",
            conversation_id="test_conv_id",
            dialogue_count=5,
        )
        pool = VariablePool()
        add_variables_to_pool(pool, sys_vars)
        json = pool.model_dump_json()
        pool2 = VariablePool.model_validate_json(json)
        assert pool2.variable_dictionary == pool.variable_dictionary

        for mode in ["json", "python"]:
            dict_ = pool.model_dump(mode=mode)
            pool2 = VariablePool.model_validate(dict_)
            assert pool2.variable_dictionary == pool.variable_dictionary

    def test_pool_without_file_vars(self):
        pool = self._create_pool_without_file()
        json = pool.model_dump_json()
        pool2 = pool.model_validate_json(json)
        assert pool2.variable_dictionary == pool.variable_dictionary
        assert pool2 == pool

    def test_basic_dictionary_round_trip(self):
        """Test basic round-trip serialization: model_dump() → model_validate()"""
        # Create a comprehensive VariablePool with all data types
        original_pool = self._create_pool_without_file()
        self._add_node_data_to_pool(original_pool)

        # Serialize to dictionary using Pydantic's model_dump()
        serialized_data = original_pool.model_dump()

        # Verify serialized data structure
        assert isinstance(serialized_data, dict)
        assert "variable_dictionary" in serialized_data

        # Deserialize back using Pydantic's model_validate()
        reconstructed_pool = VariablePool.model_validate(serialized_data)

        # Verify data integrity is preserved
        self._assert_pools_equal(original_pool, reconstructed_pool)

    def test_json_round_trip(self):
        """Test JSON round-trip serialization: model_dump_json() → model_validate_json()"""
        # Create a comprehensive VariablePool with all data types
        original_pool = self._create_pool_without_file()
        self._add_node_data_to_pool(original_pool)

        # Serialize to JSON string using Pydantic's model_dump_json()
        json_data = original_pool.model_dump_json()

        # Verify JSON is valid string
        assert isinstance(json_data, str)
        assert len(json_data) > 0

        # Deserialize back using Pydantic's model_validate_json()
        reconstructed_pool = VariablePool.model_validate_json(json_data)

        # Verify data integrity is preserved
        self._assert_pools_equal(original_pool, reconstructed_pool)

    def test_complex_data_serialization(self):
        """Test file-aware VariablePool round-trips through Dify's model boundary."""
        original_pool = self._create_pool_without_file()
        self._add_node_data_to_pool(original_pool, with_file=True)

        # Test dictionary round-trip
        dict_data = original_pool.model_dump()
        reconstructed_dict = VariablePool.model_validate(rebuild_serialized_graph_files_without_lookup(dict_data))

        # Test JSON round-trip
        json_data = original_pool.model_dump_json()
        reconstructed_json = VariablePool.model_validate(
            rebuild_serialized_graph_files_without_lookup(json.loads(json_data))
        )

        # Verify both reconstructed pools are equivalent
        self._assert_pools_equal(reconstructed_dict, reconstructed_json)
        # TODO: assert the data for file object...

    def _assert_pools_equal(self, pool1: VariablePool, pool2: VariablePool):
        """Assert that two VariablePools contain equivalent data"""

        assert pool1.variable_dictionary == pool2.variable_dictionary

        # Test key variable retrievals to ensure functionality is preserved
        test_selectors = [
            (SYSTEM_VARIABLE_NODE_ID, "user_id"),
            (SYSTEM_VARIABLE_NODE_ID, "app_id"),
            (ENVIRONMENT_VARIABLE_NODE_ID, "env_string"),
            (ENVIRONMENT_VARIABLE_NODE_ID, "env_number"),
            (CONVERSATION_VARIABLE_NODE_ID, "conv_string"),
            (self._NODE1_ID, "string_var"),
            (self._NODE1_ID, "int_var"),
            (self._NODE1_ID, "float_var"),
            (self._NODE2_ID, "array_string"),
            (self._NODE2_ID, "array_number"),
        ]

        for selector in test_selectors:
            val1 = pool1.get(selector)
            val2 = pool2.get(selector)

            # Both should exist or both should be None
            assert (val1 is None) == (val2 is None)

            if val1 is not None and val2 is not None:
                # Values should be equal
                assert val1.value == val2.value
                # Value types should be the same (more important than exact class type)
                assert val1.value_type == val2.value_type

    def test_variable_pool_deserialization_default_dict(self):
        variable_pool = VariablePool(
            variable_dictionary=defaultdict(dict),
        )
        add_variables_to_pool(variable_pool, build_system_variables(workflow_id=str(uuid.uuid4())))
        add_variables_to_pool(
            variable_pool,
            [
                StringVariable(name="str_var", value="a", selector=[ENVIRONMENT_VARIABLE_NODE_ID, "str_var"]),
                IntegerVariable(name="int_var", value=1, selector=[CONVERSATION_VARIABLE_NODE_ID, "int_var"]),
            ],
        )
        variable_pool.add(["start", "a"], 1)
        variable_pool.add(["start", "b"], "2")
        assert isinstance(variable_pool.variable_dictionary, defaultdict)
        json = variable_pool.model_dump_json()
        loaded = VariablePool.model_validate_json(json)
        assert isinstance(loaded.variable_dictionary, defaultdict)

        loaded.add(["non_exist_node", "a"], 1)

        pool_dict = variable_pool.model_dump()
        loaded = VariablePool.model_validate(pool_dict)
        assert isinstance(loaded.variable_dictionary, defaultdict)
        loaded.add(["non_exist_node", "a"], 1)


def test_get_attr():
    vp = VariablePool()
    value = {"output": StringSegment(value="hello")}

    vp.add(["node", "name"], value)
    res = vp.get(["node", "name", "output"])
    assert res is not None
    assert res.value == "hello"
