import dataclasses
from typing import Annotated

import orjson
import pytest
from pydantic import BaseModel, Discriminator, Tag

from core.helper import encrypter
from core.workflow.system_variables import build_bootstrap_variables, build_system_variables
from core.workflow.variable_pool_initializer import add_variables_to_pool
from graphon.file import File, FileTransferMethod, FileType
from graphon.runtime import VariablePool
from graphon.variables.segment_group import SegmentGroup
from graphon.variables.segments import (
    ArrayAnySegment,
    ArrayBooleanSegment,
    ArrayFileSegment,
    ArrayNumberSegment,
    ArrayObjectSegment,
    ArrayStringSegment,
    BooleanSegment,
    FileSegment,
    FloatSegment,
    IntegerSegment,
    NoneSegment,
    ObjectSegment,
    Segment,
    StringSegment,
    get_segment_discriminator,
)
from graphon.variables.types import SegmentType
from graphon.variables.utils import (
    dumps_with_segments,
    segment_orjson_default,
    to_selector,
)
from graphon.variables.variables import (
    ArrayAnyVariable,
    ArrayFileVariable,
    ArrayNumberVariable,
    ArrayObjectVariable,
    ArrayStringVariable,
    FileVariable,
    FloatVariable,
    IntegerVariable,
    NoneVariable,
    ObjectVariable,
    SecretVariable,
    StringVariable,
    Variable,
)
from models.utils.file_input_compat import rebuild_serialized_graph_files_without_lookup

type SegmentUnion = Annotated[
    (
        Annotated[NoneSegment, Tag(SegmentType.NONE)]
        | Annotated[StringSegment, Tag(SegmentType.STRING)]
        | Annotated[FloatSegment, Tag(SegmentType.FLOAT)]
        | Annotated[IntegerSegment, Tag(SegmentType.INTEGER)]
        | Annotated[ObjectSegment, Tag(SegmentType.OBJECT)]
        | Annotated[FileSegment, Tag(SegmentType.FILE)]
        | Annotated[BooleanSegment, Tag(SegmentType.BOOLEAN)]
        | Annotated[ArrayAnySegment, Tag(SegmentType.ARRAY_ANY)]
        | Annotated[ArrayStringSegment, Tag(SegmentType.ARRAY_STRING)]
        | Annotated[ArrayNumberSegment, Tag(SegmentType.ARRAY_NUMBER)]
        | Annotated[ArrayObjectSegment, Tag(SegmentType.ARRAY_OBJECT)]
        | Annotated[ArrayFileSegment, Tag(SegmentType.ARRAY_FILE)]
        | Annotated[ArrayBooleanSegment, Tag(SegmentType.ARRAY_BOOLEAN)]
    ),
    Discriminator(get_segment_discriminator),
]


def _build_variable_pool(
    *,
    system_variables: list[Variable] | None = None,
    environment_variables: list[Variable] | None = None,
) -> VariablePool:
    variable_pool = VariablePool()
    add_variables_to_pool(
        variable_pool,
        build_bootstrap_variables(
            system_variables=system_variables or [],
            environment_variables=environment_variables or [],
        ),
    )
    return variable_pool


def test_segment_group_to_text():
    variable_pool = _build_variable_pool(
        system_variables=build_system_variables(user_id="fake-user-id"),
        environment_variables=[
            SecretVariable(name="secret_key", value="fake-secret-key"),
        ],
    )
    variable_pool.add(("node_id", "custom_query"), "fake-user-query")
    template = (
        "Hello, {{#sys.user_id#}}! Your query is {{#node_id.custom_query#}}. And your key is {{#env.secret_key#}}."
    )
    segments_group = variable_pool.convert_template(template)

    assert segments_group.text == "Hello, fake-user-id! Your query is fake-user-query. And your key is fake-secret-key."
    assert segments_group.log == (
        f"Hello, fake-user-id! Your query is fake-user-query."
        f" And your key is {encrypter.obfuscated_token('fake-secret-key')}."
    )


def test_convert_constant_to_segment_group():
    variable_pool = _build_variable_pool(
        system_variables=build_system_variables(user_id="1", app_id="1", workflow_id="1"),
    )
    template = "Hello, world!"
    segments_group = variable_pool.convert_template(template)
    assert segments_group.text == "Hello, world!"
    assert segments_group.log == "Hello, world!"


def test_convert_variable_to_segment_group():
    variable_pool = _build_variable_pool(system_variables=build_system_variables(user_id="fake-user-id"))
    template = "{{#sys.user_id#}}"
    segments_group = variable_pool.convert_template(template)
    assert segments_group.text == "fake-user-id"
    assert segments_group.log == "fake-user-id"
    assert isinstance(segments_group.value[0], StringVariable)
    assert segments_group.value[0].value == "fake-user-id"


class _Segments(BaseModel):
    segments: list[SegmentUnion]


class _Variables(BaseModel):
    variables: list[Variable]


def create_test_file(
    file_type: FileType = FileType.DOCUMENT,
    transfer_method: FileTransferMethod = FileTransferMethod.LOCAL_FILE,
    filename: str = "test.txt",
    extension: str = ".txt",
    mime_type: str = "text/plain",
    size: int = 1024,
) -> File:
    """Factory function to create File objects for testing"""
    return File(
        file_type=file_type,
        transfer_method=transfer_method,
        filename=filename,
        extension=extension,
        mime_type=mime_type,
        size=size,
        related_id="test-file-id" if transfer_method != FileTransferMethod.REMOTE_URL else None,
        remote_url="https://example.com/file.txt" if transfer_method == FileTransferMethod.REMOTE_URL else None,
        storage_key="test-storage-key",
    )


class TestSegmentDumpAndLoad:
    """Test suite for segment and variable serialization/deserialization"""

    def test_segments(self):
        """Test basic segment serialization compatibility"""
        model = _Segments(segments=[IntegerSegment(value=1), StringSegment(value="a")])
        json = model.model_dump_json()
        loaded = _Segments.model_validate_json(json)
        assert loaded == model

    def test_segment_number(self):
        """Test number segment serialization compatibility"""
        model = _Segments(segments=[IntegerSegment(value=1), FloatSegment(value=1.0)])
        json = model.model_dump_json()
        loaded = _Segments.model_validate_json(json)
        assert loaded == model

    def test_variables(self):
        """Test variable serialization compatibility"""
        model = _Variables(variables=[IntegerVariable(value=1, name="int"), StringVariable(value="a", name="str")])
        json = model.model_dump_json()
        restored = _Variables.model_validate_json(json)
        assert restored == model

    def test_all_segments_serialization(self):
        """Test file-aware segment serialization through Dify's model boundary."""
        # Create one instance of each segment type
        test_file = create_test_file()

        all_segments: list[SegmentUnion] = [
            NoneSegment(),
            StringSegment(value="test string"),
            IntegerSegment(value=42),
            FloatSegment(value=3.14),
            ObjectSegment(value={"key": "value", "number": 123}),
            FileSegment(value=test_file),
            ArrayAnySegment(value=[1, "string", 3.14, {"key": "value"}]),
            ArrayStringSegment(value=["hello", "world"]),
            ArrayNumberSegment(value=[1, 2.5, 3]),
            ArrayObjectSegment(value=[{"id": 1}, {"id": 2}]),
            ArrayFileSegment(value=[]),  # Empty array to avoid file complexity
        ]

        # Test serialization and deserialization
        model = _Segments(segments=all_segments)
        json_str = model.model_dump_json()
        loaded = _Segments.model_validate(rebuild_serialized_graph_files_without_lookup(orjson.loads(json_str)))

        # Verify all segments are preserved
        assert len(loaded.segments) == len(all_segments)

        for original, loaded_segment in zip(all_segments, loaded.segments):
            assert type(loaded_segment) == type(original)
            assert loaded_segment.value_type == original.value_type

            # For file segments, compare key properties instead of exact equality
            if isinstance(original, FileSegment) and isinstance(loaded_segment, FileSegment):
                orig_file = original.value
                loaded_file = loaded_segment.value
                assert isinstance(orig_file, File)
                assert isinstance(loaded_file, File)
                assert loaded_file.type == orig_file.type
                assert loaded_file.filename == orig_file.filename
            else:
                assert loaded_segment.value == original.value

    def test_all_variables_serialization(self):
        """Test file-aware variable serialization through Dify's model boundary."""
        # Create one instance of each variable type
        test_file = create_test_file()

        all_variables: list[Variable] = [
            NoneVariable(name="none_var"),
            StringVariable(value="test string", name="string_var"),
            IntegerVariable(value=42, name="int_var"),
            FloatVariable(value=3.14, name="float_var"),
            ObjectVariable(value={"key": "value", "number": 123}, name="object_var"),
            FileVariable(value=test_file, name="file_var"),
            ArrayAnyVariable(value=[1, "string", 3.14, {"key": "value"}], name="array_any_var"),
            ArrayStringVariable(value=["hello", "world"], name="array_string_var"),
            ArrayNumberVariable(value=[1, 2.5, 3], name="array_number_var"),
            ArrayObjectVariable(value=[{"id": 1}, {"id": 2}], name="array_object_var"),
            ArrayFileVariable(value=[], name="array_file_var"),  # Empty array to avoid file complexity
        ]

        # Test serialization and deserialization
        model = _Variables(variables=all_variables)
        json_str = model.model_dump_json()
        loaded = _Variables.model_validate(rebuild_serialized_graph_files_without_lookup(orjson.loads(json_str)))

        # Verify all variables are preserved
        assert len(loaded.variables) == len(all_variables)

        for original, loaded_variable in zip(all_variables, loaded.variables):
            assert type(loaded_variable) == type(original)
            assert loaded_variable.value_type == original.value_type
            assert loaded_variable.name == original.name

            # For file variables, compare key properties instead of exact equality
            if isinstance(original, FileVariable) and isinstance(loaded_variable, FileVariable):
                orig_file = original.value
                loaded_file = loaded_variable.value
                assert isinstance(orig_file, File)
                assert isinstance(loaded_file, File)
                assert loaded_file.type == orig_file.type
                assert loaded_file.filename == orig_file.filename
            else:
                assert loaded_variable.value == original.value

    def test_segment_discriminator_function_for_segment_types(self):
        """Test the segment discriminator function"""

        @dataclasses.dataclass
        class TestCase:
            segment: Segment
            expected_segment_type: SegmentType

        file1 = create_test_file()
        file2 = create_test_file(filename="test2.txt")

        cases = [
            TestCase(
                NoneSegment(),
                SegmentType.NONE,
            ),
            TestCase(
                StringSegment(value=""),
                SegmentType.STRING,
            ),
            TestCase(
                FloatSegment(value=0.0),
                SegmentType.FLOAT,
            ),
            TestCase(
                IntegerSegment(value=0),
                SegmentType.INTEGER,
            ),
            TestCase(
                ObjectSegment(value={}),
                SegmentType.OBJECT,
            ),
            TestCase(
                FileSegment(value=file1),
                SegmentType.FILE,
            ),
            TestCase(
                ArrayAnySegment(value=[0, 0.0, ""]),
                SegmentType.ARRAY_ANY,
            ),
            TestCase(
                ArrayStringSegment(value=[""]),
                SegmentType.ARRAY_STRING,
            ),
            TestCase(
                ArrayNumberSegment(value=[0, 0.0]),
                SegmentType.ARRAY_NUMBER,
            ),
            TestCase(
                ArrayObjectSegment(value=[{}]),
                SegmentType.ARRAY_OBJECT,
            ),
            TestCase(
                ArrayFileSegment(value=[file1, file2]),
                SegmentType.ARRAY_FILE,
            ),
        ]

        for test_case in cases:
            segment = test_case.segment
            assert get_segment_discriminator(segment) == test_case.expected_segment_type, (
                f"get_segment_discriminator failed for type {type(segment)}"
            )
            model_dict = segment.model_dump(mode="json")
            assert get_segment_discriminator(model_dict) == test_case.expected_segment_type, (
                f"get_segment_discriminator failed for serialized form of type {type(segment)}"
            )

    def test_variable_discriminator_function_for_variable_types(self):
        """Test the variable discriminator function"""

        @dataclasses.dataclass
        class TestCase:
            variable: Variable
            expected_segment_type: SegmentType

        file1 = create_test_file()
        file2 = create_test_file(filename="test2.txt")

        cases = [
            TestCase(
                NoneVariable(name="none_var"),
                SegmentType.NONE,
            ),
            TestCase(
                StringVariable(value="test", name="string_var"),
                SegmentType.STRING,
            ),
            TestCase(
                FloatVariable(value=0.0, name="float_var"),
                SegmentType.FLOAT,
            ),
            TestCase(
                IntegerVariable(value=0, name="int_var"),
                SegmentType.INTEGER,
            ),
            TestCase(
                ObjectVariable(value={}, name="object_var"),
                SegmentType.OBJECT,
            ),
            TestCase(
                FileVariable(value=file1, name="file_var"),
                SegmentType.FILE,
            ),
            TestCase(
                SecretVariable(value="secret", name="secret_var"),
                SegmentType.SECRET,
            ),
            TestCase(
                ArrayAnyVariable(value=[0, 0.0, ""], name="array_any_var"),
                SegmentType.ARRAY_ANY,
            ),
            TestCase(
                ArrayStringVariable(value=[""], name="array_string_var"),
                SegmentType.ARRAY_STRING,
            ),
            TestCase(
                ArrayNumberVariable(value=[0, 0.0], name="array_number_var"),
                SegmentType.ARRAY_NUMBER,
            ),
            TestCase(
                ArrayObjectVariable(value=[{}], name="array_object_var"),
                SegmentType.ARRAY_OBJECT,
            ),
            TestCase(
                ArrayFileVariable(value=[file1, file2], name="array_file_var"),
                SegmentType.ARRAY_FILE,
            ),
        ]

        for test_case in cases:
            variable = test_case.variable
            assert get_segment_discriminator(variable) == test_case.expected_segment_type, (
                f"get_segment_discriminator failed for type {type(variable)}"
            )
            model_dict = variable.model_dump(mode="json")
            assert get_segment_discriminator(model_dict) == test_case.expected_segment_type, (
                f"get_segment_discriminator failed for serialized form of type {type(variable)}"
            )

    def test_invalid_value_for_discriminator(self):
        # Test invalid cases
        assert get_segment_discriminator({"value_type": "invalid"}) is None
        assert get_segment_discriminator({}) is None
        assert get_segment_discriminator("not_a_dict") is None
        assert get_segment_discriminator(42) is None
        assert get_segment_discriminator(object) is None


class TestSegmentAdditionalProperties:
    def test_base_segment_text_log_markdown_size_and_to_object(self):
        """Ensure StringSegment exposes text, log, markdown, size and to_object."""
        segment = StringSegment(value="hello")

        assert segment.text == "hello"
        assert segment.log == "hello"
        assert segment.markdown == "hello"
        assert segment.size > 0
        assert segment.to_object() == "hello"

    def test_none_segment_empty_outputs(self):
        """Ensure NoneSegment renders empty text, log and markdown."""
        segment = NoneSegment()

        assert segment.text == ""
        assert segment.log == ""
        assert segment.markdown == ""

    def test_object_segment_json_outputs(self):
        """Ensure ObjectSegment renders JSON output for text, log and markdown."""
        segment = ObjectSegment(value={"key": "值", "n": 1})

        assert segment.text == '{"key": "值", "n": 1}'
        assert segment.log == '{\n  "key": "值",\n  "n": 1\n}'
        assert segment.markdown == '{\n  "key": "值",\n  "n": 1\n}'

    def test_array_segment_text_and_markdown(self):
        """Ensure ArrayAnySegment handles empty/non-empty text and markdown rendering."""
        empty_segment = ArrayAnySegment(value=[])
        non_empty_segment = ArrayAnySegment(value=[1, "two"])

        assert empty_segment.text == ""
        assert non_empty_segment.text == "[1, 'two']"
        assert non_empty_segment.markdown == "- 1\n- two"

    def test_file_segment_properties(self):
        """Ensure FileSegment markdown, text and log fields match expected behavior."""
        file = create_test_file(transfer_method=FileTransferMethod.REMOTE_URL, filename="doc.txt")
        segment = FileSegment(value=file)

        assert segment.markdown == "[doc.txt](https://example.com/file.txt)"
        assert segment.log == ""
        assert segment.text == ""

    def test_array_string_segment_text_branches(self):
        """Ensure ArrayStringSegment text handling for empty and non-empty values."""
        empty_segment = ArrayStringSegment(value=[])
        non_empty_segment = ArrayStringSegment(value=["hello", "世界"])

        assert empty_segment.text == ""
        assert non_empty_segment.text == '["hello", "世界"]'

    def test_array_file_segment_markdown_and_empty_text_log(self):
        """Ensure ArrayFileSegment markdown renders links and text/log stay empty."""
        file1 = create_test_file(transfer_method=FileTransferMethod.REMOTE_URL, filename="a.txt")
        file2 = create_test_file(transfer_method=FileTransferMethod.REMOTE_URL, filename="b.txt")
        segment = ArrayFileSegment(value=[file1, file2])

        assert segment.markdown == "[a.txt](https://example.com/file.txt)\n[b.txt](https://example.com/file.txt)"
        assert segment.log == ""
        assert segment.text == ""


class TestSegmentGroupAdditional:
    def test_segment_group_markdown_and_to_object(self):
        group = SegmentGroup(value=[StringSegment(value="A"), NoneSegment(), StringSegment(value="B")])

        assert group.markdown == "AB"
        assert group.to_object() == ["A", None, "B"]


class TestSegmentUtils:
    def test_to_selector_without_paths(self):
        assert to_selector("node-1", "output") == ["node-1", "output"]

    def test_to_selector_with_paths(self):
        assert to_selector("node-1", "output", ("a", "b")) == ["node-1", "output", "a", "b"]

    def test_array_file_segment_serialization(self):
        file1 = create_test_file(transfer_method=FileTransferMethod.REMOTE_URL, filename="a.txt")
        file2 = create_test_file(transfer_method=FileTransferMethod.REMOTE_URL, filename="b.txt")

        result = segment_orjson_default(ArrayFileSegment(value=[file1, file2]))

        assert len(result) == 2
        assert result[0]["filename"] == "a.txt"
        assert result[1]["filename"] == "b.txt"

    def test_file_segment_serialization(self):
        file = create_test_file(transfer_method=FileTransferMethod.REMOTE_URL, filename="single.txt")

        result = segment_orjson_default(FileSegment(value=file))

        assert result["filename"] == "single.txt"
        assert result["remote_url"] == "https://example.com/file.txt"

    def test_segment_group_and_segment_serialization(self):
        group = SegmentGroup(value=[StringSegment(value="a"), StringSegment(value="b")])

        assert segment_orjson_default(group) == ["a", "b"]
        assert segment_orjson_default(StringSegment(value="value")) == "value"

    def test_segment_orjson_default_unsupported_type(self):
        with pytest.raises(TypeError, match="not JSON serializable"):
            segment_orjson_default(object())

    def test_dumps_with_segments(self):
        data = {
            "segment": StringSegment(value="hello"),
            "group": SegmentGroup(value=[StringSegment(value="x"), StringSegment(value="y")]),
            1: "numeric-key",
        }

        dumped = dumps_with_segments(data)
        loaded = orjson.loads(dumped)

        assert loaded["segment"] == "hello"
        assert loaded["group"] == ["x", "y"]
        assert loaded["1"] == "numeric-key"
