#  Copyright (c) "Neo4j"
#  Neo4j Sweden AB [https://neo4j.com]
#  #
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#  #
#      https://www.apache.org/licenses/LICENSE-2.0
#  #
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
from __future__ import annotations

import json
from typing import Tuple, Any
from unittest.mock import AsyncMock, patch, Mock

import pytest
from pydantic import ValidationError

from neo4j_graphrag.exceptions import SchemaValidationError, SchemaExtractionError
from neo4j_graphrag.experimental.components.schema import (
    SchemaBuilder,
    NodeType,
    PropertyType,
    RelationshipType,
    ConstraintType,
    GraphConstraintType,
    SchemaFromTextExtractor,
    GraphSchema,
    SchemaFromExistingGraphExtractor,
    Pattern,
)
import os
import tempfile
import yaml

from neo4j_graphrag.generation import PromptTemplate
from neo4j_graphrag.llm.types import LLMResponse
from neo4j_graphrag.utils.file_handler import FileFormat

# Most tests below read ``PropertyType.required`` to assert migration; silence only that
# deprecation. Explicit warning behavior is covered in ``test_property_type_deprecation.py``.
pytestmark = pytest.mark.filterwarnings(
    "ignore:Use GraphSchema.constraints with type EXISTENCE:DeprecationWarning"
)


def test_node_type_initialization_from_string() -> None:
    node_type = NodeType.model_validate("Label")
    assert isinstance(node_type, NodeType)
    assert node_type.label == "Label"
    assert len(node_type.properties) == 1
    assert node_type.properties[0].name == "name"
    assert node_type.properties[0].type == "STRING"


def test_node_type_additional_properties_default() -> None:
    # default behavior:
    node_type = NodeType.model_validate({"label": "Label"})
    assert node_type.additional_properties is True
    # Empty properties list now violates min_length=1 constraint
    with pytest.raises(ValidationError):
        NodeType.model_validate({"label": "Label", "properties": []})
    node_type = NodeType.model_validate(
        {"label": "Label", "properties": [{"name": "name", "type": "STRING"}]}
    )
    assert node_type.additional_properties is False

    # manually changing the default value
    # impossible cases: no properties and no additional
    with pytest.raises(ValidationError):
        NodeType.model_validate({"label": "Label", "additional_properties": False})
    with pytest.raises(ValidationError):
        NodeType.model_validate(
            {"label": "Label", "properties": [], "additional_properties": False}
        )

    # working case: properties and additional allowed
    node_type = NodeType.model_validate(
        {
            "label": "Label",
            "properties": [{"name": "name", "type": "STRING"}],
            "additional_properties": True,
        }
    )
    assert node_type.additional_properties is True


@pytest.mark.parametrize(
    "label",
    ["__Person", "Person__", "__Person__", "__KNOWS", "KNOWS__", "__KNOWS__"],
)
def test_node_type_rejects_dunder_label(label: str) -> None:
    with pytest.raises(ValidationError, match="reserved '__'"):
        NodeType.model_validate(
            {"label": label, "properties": [{"name": "name", "type": "STRING"}]}
        )


@pytest.mark.parametrize(
    "label",
    ["__KNOWS", "KNOWS__", "__KNOWS__"],
)
def test_relationship_type_rejects_dunder_label(label: str) -> None:
    with pytest.raises(ValidationError, match="reserved '__'"):
        RelationshipType.model_validate({"label": label})


def test_property_type_initalization() -> None:
    prop = PropertyType(name="email", type="STRING")
    assert prop.name == "email"
    assert prop.type == "STRING"
    assert prop.required is False


def test_property_type_with_required_true() -> None:
    prop = PropertyType(name="id", type="INTEGER", required=True)
    assert prop.required is True


def test_property_type_is_frozen() -> None:
    prop = PropertyType(name="email", type="STRING", required=False)

    with pytest.raises(ValidationError):
        prop.name = "other"

    with pytest.raises(ValidationError):
        prop.required = True


def test_relationship_type_initialization_from_string() -> None:
    relationship_type = RelationshipType.model_validate("REL")
    assert isinstance(relationship_type, RelationshipType)
    assert relationship_type.label == "REL"
    assert relationship_type.properties == []


def test_relationship_type_additional_properties_default() -> None:
    relationship_type = RelationshipType.model_validate({"label": "REL"})
    assert relationship_type.additional_properties is True
    relationship_type = RelationshipType.model_validate(
        {"label": "REL", "properties": []}
    )
    assert relationship_type.additional_properties is True
    relationship_type = RelationshipType.model_validate(
        {"label": "REL", "properties": [{"name": "name", "type": "STRING"}]}
    )
    assert relationship_type.additional_properties is False

    # manually changing the default value
    # auto-correction: no properties and additional_properties=False -> auto-corrected to True
    relationship_type = RelationshipType.model_validate(
        {"label": "REL", "additional_properties": False}
    )
    assert relationship_type.additional_properties is True  # auto-corrected
    relationship_type = RelationshipType.model_validate(
        {"label": "REL", "properties": [], "additional_properties": False}
    )
    assert relationship_type.additional_properties is True  # auto-corrected

    # working case: properties and additional allowed
    relationship_type = RelationshipType.model_validate(
        {
            "label": "REL",
            "properties": [{"name": "name", "type": "STRING"}],
            "additional_properties": True,
        }
    )
    assert relationship_type.additional_properties is True


def test_constraint_type_initialization() -> None:
    constraint = ConstraintType(
        type=GraphConstraintType.UNIQUENESS,
        node_type="Person",
        property_names=("name",),
    )
    assert constraint.type == "UNIQUENESS"
    assert constraint.node_type == "Person"
    assert constraint.property_names == ("name",)


def test_constraint_type_is_frozen() -> None:
    constraint = ConstraintType(
        type=GraphConstraintType.UNIQUENESS,
        node_type="Person",
        property_names=("name",),
    )

    with pytest.raises(ValidationError):
        constraint.type = GraphConstraintType.UNIQUENESS

    with pytest.raises(ValidationError):
        constraint.node_type = "Organization"

    with pytest.raises(ValidationError):
        constraint.property_names = ("id",)


def test_constraint_type_rejects_empty_property_names() -> None:
    """``property_names`` has ``min_length=1``; empty tuple/list fails validation."""
    with pytest.raises(ValidationError) as exc_info:
        ConstraintType(
            type=GraphConstraintType.UNIQUENESS,
            node_type="Person",
            property_names=(),
        )
    assert "too_short" in str(exc_info.value).lower()

    with pytest.raises(ValidationError):
        ConstraintType.model_validate(
            {
                "type": "UNIQUENESS",
                "node_type": "Person",
                "property_names": [],
            }
        )


def test_schema_additional_node_types_default() -> None:
    schema_dict: dict[str, Any] = {
        "node_types": [],
    }
    schema = GraphSchema.model_validate(schema_dict)
    assert schema.additional_node_types is True

    schema_dict = {
        "node_types": ["Person"],
    }
    schema = GraphSchema.model_validate(schema_dict)
    assert schema.additional_node_types is False


def test_schema_additional_relationship_types_default() -> None:
    schema_dict: dict[str, Any] = {
        "node_types": [],
    }
    schema = GraphSchema.model_validate(schema_dict)
    assert schema.additional_relationship_types is True

    schema_dict = {
        "node_types": [],
        "relationship_types": ["REL"],
    }
    schema = GraphSchema.model_validate(schema_dict)
    assert schema.additional_relationship_types is False


def test_schema_additional_patterns_default() -> None:
    schema_dict: dict[str, Any] = {
        "node_types": [],
    }
    schema = GraphSchema.model_validate(schema_dict)
    assert schema.additional_patterns is True

    schema_dict = {
        "node_types": ["Person"],
        "relationship_types": ["REL"],
        "patterns": [("Person", "REL", "Person")],
    }
    schema = GraphSchema.model_validate(schema_dict)
    assert schema.additional_patterns is False


def test_schema_additional_parameter_validation() -> None:
    """Additional relationship types not allowed, but additional patterns allowed

    => raise Exception
    """
    schema_dict = {
        "node_types": [
            {
                "label": "Person",
                "properties": [
                    {
                        "name": "name",
                        "type": "STRING",
                    },
                    {"name": "height", "type": "INTEGER"},
                ],
            }
        ],
        "relationship_types": [
            {
                "label": "KNOWS",
            }
        ],
        "patterns": [
            ("Person", "KNOWS", "Person"),
        ],
        "additional_relationship_types": True,
        "additional_patterns": False,
    }
    with pytest.raises(
        ValidationError,
        match="`additional_relationship_types` must be set to False when using `additional_patterns=False`",
    ):
        GraphSchema.model_validate(schema_dict)


def test_schema_constraint_validation_property_not_in_node_type() -> None:
    schema_dict: dict[str, Any] = {
        "node_types": [
            {
                "label": "Person",
                "properties": [{"name": "name", "type": "STRING"}],
            }
        ],
        "constraints": [
            {"type": "UNIQUENESS", "node_type": "Person", "property_name": "email"}
        ],
    }

    with pytest.raises(SchemaValidationError) as exc_info:
        GraphSchema.model_validate(schema_dict)

    assert "Constraint references undefined property" in str(exc_info.value)
    assert "on node type 'Person'" in str(exc_info.value)


def test_schema_constraint_with_additional_properties_with_allows_unknown_property() -> (
    None
):
    # if additional_properties is True, we can define constraints that are not in the node_type
    schema_dict: dict[str, Any] = {
        "node_types": [
            {
                "label": "Person",
                "properties": [{"name": "name", "type": "STRING"}],
                "additional_properties": True,
            }
        ],
        "constraints": [
            {"type": "UNIQUENESS", "node_type": "Person", "property_name": "email"}
        ],
    }

    # Should raise - email is not allowed because the property is not defined in the node
    with pytest.raises(SchemaValidationError) as exc_info:
        GraphSchema.model_validate(schema_dict)

    assert "Constraint references undefined property 'email'" in str(exc_info.value)


def test_schema_with_valid_constraints() -> None:
    schema_dict: dict[str, Any] = {
        "node_types": [
            {"label": "Person", "properties": [{"name": "name", "type": "STRING"}]}
        ],
        "constraints": [
            {"type": "UNIQUENESS", "node_type": "Person", "property_name": "name"}
        ],
    }
    schema = GraphSchema.model_validate(schema_dict)

    assert len(schema.constraints) == 1
    assert schema.constraints[0].type == "UNIQUENESS"
    assert schema.constraints[0].node_type == "Person"
    assert schema.constraints[0].property_names == ("name",)


def test_schema_constraint_validation_invalid_node_type() -> None:
    schema_dict: dict[str, Any] = {
        "node_types": [
            {"label": "Person", "properties": [{"name": "name", "type": "STRING"}]}
        ],
        "constraints": [
            {
                "type": "UNIQUENESS",
                "node_type": "NonExistentNode",
                "property_name": "id",
            }
        ],
    }

    with pytest.raises(SchemaValidationError) as exc_info:
        GraphSchema.model_validate(schema_dict)

    assert "Constraint references undefined node type: NonExistentNode" in str(
        exc_info.value
    )


def test_schema_constraint_validation_missing_property_name() -> None:
    schema_dict: dict[str, Any] = {
        "node_types": [
            {"label": "Person", "properties": [{"name": "name", "type": "STRING"}]}
        ],
        "constraints": [
            {"type": "UNIQUENESS", "node_type": "Person", "property_name": ""}
        ],
    }

    with pytest.raises((SchemaValidationError, ValidationError)):
        GraphSchema.model_validate(schema_dict)


def test_schema_key_constraint_node_valid() -> None:
    schema_dict: dict[str, Any] = {
        "node_types": [
            {"label": "Person", "properties": [{"name": "email", "type": "STRING"}]}
        ],
        "constraints": [
            {
                "type": "KEY",
                "node_type": "Person",
                "property_name": "email",
                "relationship_type": None,
            }
        ],
    }
    schema = GraphSchema.model_validate(schema_dict)
    assert schema.key_property_names_for_node("Person") == {"email"}
    assert schema.mandatory_property_names_for_node("Person") == {"email"}
    assert schema.uniqueness_property_names_for_node("Person") == set()


def test_schema_key_constraint_relationship_valid() -> None:
    schema_dict: dict[str, Any] = {
        "node_types": [
            {"label": "Person", "properties": [{"name": "name", "type": "STRING"}]}
        ],
        "relationship_types": [
            {
                "label": "KNOWS",
                "properties": [{"name": "since", "type": "INTEGER"}],
            }
        ],
        "constraints": [
            {
                "type": "KEY",
                "node_type": "",
                "property_name": "since",
                "relationship_type": "KNOWS",
            }
        ],
    }
    schema = GraphSchema.model_validate(schema_dict)
    assert schema.key_property_names_for_relationship("KNOWS") == {"since"}
    assert schema.mandatory_property_names_for_relationship("KNOWS") == {"since"}


def test_schema_uniqueness_and_key_same_property_rejected() -> None:
    schema_dict: dict[str, Any] = {
        "node_types": [
            {"label": "Person", "properties": [{"name": "id", "type": "STRING"}]}
        ],
        "constraints": [
            {
                "type": "UNIQUENESS",
                "node_type": "Person",
                "property_name": "id",
                "relationship_type": None,
            },
            {
                "type": "KEY",
                "node_type": "Person",
                "property_name": "id",
                "relationship_type": None,
            },
        ],
    }
    with pytest.raises(SchemaValidationError, match="UNIQUENESS and KEY"):
        GraphSchema.model_validate(schema_dict)


# --- Composite (multi-property) constraint tests ---


def test_composite_key_constraint_node_valid() -> None:
    schema_dict: dict[str, Any] = {
        "node_types": [
            {
                "label": "Actor",
                "properties": [
                    {"name": "firstname", "type": "STRING"},
                    {"name": "surname", "type": "STRING"},
                ],
            }
        ],
        "constraints": [
            {
                "type": "KEY",
                "node_type": "Actor",
                "property_names": ["firstname", "surname"],
                "relationship_type": None,
            }
        ],
    }
    schema = GraphSchema.model_validate(schema_dict)
    assert len(schema.constraints) == 1
    assert schema.constraints[0].property_names == ("firstname", "surname")
    assert schema.key_property_names_for_node("Actor") == {"firstname", "surname"}
    assert schema.mandatory_property_names_for_node("Actor") == {
        "firstname",
        "surname",
    }


def test_composite_uniqueness_constraint_valid() -> None:
    schema_dict: dict[str, Any] = {
        "node_types": [
            {
                "label": "Book",
                "properties": [
                    {"name": "title", "type": "STRING"},
                    {"name": "year", "type": "INTEGER"},
                ],
            }
        ],
        "constraints": [
            {
                "type": "UNIQUENESS",
                "node_type": "Book",
                "property_names": ["title", "year"],
            }
        ],
    }
    schema = GraphSchema.model_validate(schema_dict)
    assert len(schema.constraints) == 1
    assert schema.constraints[0].property_names == ("title", "year")
    assert schema.uniqueness_property_names_for_node("Book") == {"title", "year"}


def test_existence_rejects_composite() -> None:
    schema_dict: dict[str, Any] = {
        "node_types": [
            {
                "label": "Person",
                "properties": [
                    {"name": "name", "type": "STRING"},
                    {"name": "email", "type": "STRING"},
                ],
            }
        ],
        "constraints": [
            {
                "type": "EXISTENCE",
                "node_type": "Person",
                "property_names": ["name", "email"],
            }
        ],
    }
    with pytest.raises(ValueError, match="EXISTENCE constraint must have exactly one"):
        GraphSchema.model_validate(schema_dict)


def test_backward_compat_property_name_migrates_to_property_names() -> None:
    schema_dict: dict[str, Any] = {
        "node_types": [
            {"label": "Person", "properties": [{"name": "name", "type": "STRING"}]}
        ],
        "constraints": [
            {
                "type": "UNIQUENESS",
                "node_type": "Person",
                "property_name": "name",
            }
        ],
    }
    schema = GraphSchema.model_validate(schema_dict)
    assert schema.constraints[0].property_names == ("name",)
    with pytest.warns(DeprecationWarning, match="property_names"):
        assert schema.constraints[0].property_name == "name"


def test_composite_constraint_validates_all_properties_exist() -> None:
    schema_dict: dict[str, Any] = {
        "node_types": [
            {
                "label": "Actor",
                "properties": [
                    {"name": "firstname", "type": "STRING"},
                ],
            }
        ],
        "constraints": [
            {
                "type": "KEY",
                "node_type": "Actor",
                "property_names": ["firstname", "nonexistent"],
                "relationship_type": None,
            }
        ],
    }
    with pytest.raises(SchemaValidationError, match="nonexistent"):
        GraphSchema.model_validate(schema_dict)


def test_composite_uniqueness_and_key_same_properties_rejected() -> None:
    schema_dict: dict[str, Any] = {
        "node_types": [
            {
                "label": "Actor",
                "properties": [
                    {"name": "firstname", "type": "STRING"},
                    {"name": "surname", "type": "STRING"},
                ],
            }
        ],
        "constraints": [
            {
                "type": "UNIQUENESS",
                "node_type": "Actor",
                "property_names": ["firstname", "surname"],
            },
            {
                "type": "KEY",
                "node_type": "Actor",
                "property_names": ["firstname", "surname"],
                "relationship_type": None,
            },
        ],
    }
    with pytest.raises(SchemaValidationError, match="UNIQUENESS and KEY"):
        GraphSchema.model_validate(schema_dict)


def test_composite_uniqueness_and_key_different_properties_allowed() -> None:
    schema_dict: dict[str, Any] = {
        "node_types": [
            {
                "label": "Actor",
                "properties": [
                    {"name": "firstname", "type": "STRING"},
                    {"name": "surname", "type": "STRING"},
                    {"name": "email", "type": "STRING"},
                ],
            }
        ],
        "constraints": [
            {
                "type": "UNIQUENESS",
                "node_type": "Actor",
                "property_names": ["email"],
            },
            {
                "type": "KEY",
                "node_type": "Actor",
                "property_names": ["firstname", "surname"],
                "relationship_type": None,
            },
        ],
    }
    schema = GraphSchema.model_validate(schema_dict)
    assert len(schema.constraints) == 2
    assert schema.key_property_names_for_node("Actor") == {"firstname", "surname"}
    assert schema.uniqueness_property_names_for_node("Actor") == {"email"}


def test_extract_graph_constraints_from_metadata_composite_key() -> None:
    from neo4j_graphrag.experimental.components.schema import (
        SchemaFromExistingGraphExtractor,
    )

    structured_schema: dict[str, Any] = {
        "metadata": {
            "constraint": [
                {
                    "type": "NODE_KEY",
                    "properties": ["firstname", "surname"],
                    "labelsOrTypes": ["Actor"],
                }
            ]
        }
    }
    result = SchemaFromExistingGraphExtractor._extract_graph_constraints_from_metadata(
        structured_schema
    )
    assert len(result) == 1
    assert result[0]["type"] == "KEY"
    assert result[0]["node_type"] == "Actor"
    assert result[0]["property_names"] == ("firstname", "surname")
    assert result[0]["property_name"] == "firstname"


@pytest.fixture
def valid_node_types() -> tuple[NodeType, ...]:
    # required=False so tuples match GraphSchema after validation (legacy required=True
    # is migrated to EXISTENCE and cleared).
    return (
        NodeType(
            label="PERSON",
            description="An individual human being.",
            properties=[
                PropertyType(name="birth date", type="ZONED_DATETIME"),
                PropertyType(name="name", type="STRING", required=False),
            ],
            additional_properties=False,
        ),
        NodeType(
            label="ORGANIZATION",
            description="A structured group of people with a common purpose.",
            properties=[PropertyType(name="name", type="STRING")],
            additional_properties=True,
        ),
        NodeType(
            label="AGE",
            description="Age of a person in years.",
            properties=[PropertyType(name="value", type="INTEGER")],
            additional_properties=True,
        ),
    )


@pytest.fixture
def valid_relationship_types() -> tuple[RelationshipType, ...]:
    return (
        RelationshipType(
            label="EMPLOYED_BY",
            description="Indicates employment relationship.",
            properties=[
                PropertyType(name="start_time", type="LOCAL_DATETIME", required=False),
                PropertyType(name="end_time", type="LOCAL_DATETIME"),
            ],
            additional_properties=False,
        ),
        RelationshipType(
            label="ORGANIZED_BY",
            description="Indicates organization responsible for an event.",
        ),
        RelationshipType(
            label="ATTENDED_BY", description="Indicates attendance at an event."
        ),
    )


@pytest.fixture
def valid_patterns() -> tuple[Pattern, ...]:
    return (
        Pattern(source="PERSON", relationship="EMPLOYED_BY", target="ORGANIZATION"),
        Pattern(source="ORGANIZATION", relationship="ATTENDED_BY", target="PERSON"),
    )


@pytest.fixture
def patterns_with_invalid_entity() -> tuple[Pattern, ...]:
    return (
        Pattern(source="PERSON", relationship="EMPLOYED_BY", target="ORGANIZATION"),
        Pattern(
            source="NON_EXISTENT_ENTITY", relationship="ATTENDED_BY", target="PERSON"
        ),
    )


@pytest.fixture
def valid_constraints() -> tuple[ConstraintType, ...]:
    return (
        ConstraintType(
            type=GraphConstraintType.UNIQUENESS,
            node_type="PERSON",
            property_names=("name",),
        ),
    )


@pytest.fixture
def patterns_with_invalid_relation() -> tuple[Pattern, ...]:
    return (
        Pattern(
            source="PERSON", relationship="NON_EXISTENT_RELATION", target="ORGANIZATION"
        ),
    )


@pytest.fixture
def schema_builder() -> SchemaBuilder:
    return SchemaBuilder()


@pytest.fixture
def graph_schema(
    schema_builder: SchemaBuilder,
    valid_node_types: Tuple[NodeType, ...],
    valid_relationship_types: Tuple[RelationshipType, ...],
    valid_patterns: tuple[Pattern, ...],
) -> GraphSchema:
    return schema_builder.create_schema_model(
        list(valid_node_types), list(valid_relationship_types), list(valid_patterns)
    )


def test_create_schema_model_valid_data(
    schema_builder: SchemaBuilder,
    valid_node_types: Tuple[NodeType, ...],
    valid_relationship_types: Tuple[RelationshipType, ...],
    valid_patterns: tuple[Pattern, ...],
) -> None:
    schema = schema_builder.create_schema_model(
        list(valid_node_types), list(valid_relationship_types), list(valid_patterns)
    )

    assert schema.node_types == valid_node_types
    assert schema.relationship_types == valid_relationship_types
    assert schema.patterns == valid_patterns
    assert schema.additional_node_types is False
    assert schema.additional_relationship_types is False
    assert schema.additional_patterns is False


def test_create_schema_model_with_constraints(
    schema_builder: SchemaBuilder,
    valid_node_types: Tuple[NodeType, ...],
    valid_constraints: Tuple[ConstraintType, ...],
) -> None:
    schema = schema_builder.create_schema_model(
        list(valid_node_types),
        constraints=list(valid_constraints),
    )

    assert schema.node_types == valid_node_types
    assert schema.constraints == valid_constraints
    assert len(schema.constraints) == 1
    assert schema.constraints[0].type == "UNIQUENESS"
    assert schema.constraints[0].node_type == "PERSON"
    assert schema.constraints[0].property_names == ("name",)


@pytest.mark.asyncio
async def test_run_method(
    schema_builder: SchemaBuilder,
    valid_node_types: Tuple[NodeType, ...],
    valid_relationship_types: Tuple[RelationshipType, ...],
    valid_patterns: tuple[Pattern, ...],
) -> None:
    with patch.object(
        schema_builder,
        "create_schema_model",
        return_value=GraphSchema(
            node_types=valid_node_types,
            relationship_types=valid_relationship_types,
            patterns=valid_patterns,
        ),
    ):
        schema = await schema_builder.run(
            list(valid_node_types), list(valid_relationship_types), list(valid_patterns)
        )

    assert schema.node_types == valid_node_types
    assert schema.relationship_types == valid_relationship_types
    assert schema.patterns == valid_patterns
    assert schema.additional_node_types is False
    assert schema.additional_relationship_types is False
    assert schema.additional_patterns is False


@pytest.mark.asyncio
async def test_run_method_with_constraints(
    schema_builder: SchemaBuilder,
    valid_node_types: Tuple[NodeType, ...],
    valid_constraints: Tuple[ConstraintType, ...],
) -> None:
    schema = await schema_builder.run(
        list(valid_node_types),
        constraints=list(valid_constraints),
    )

    assert schema.node_types == valid_node_types
    assert schema.constraints == valid_constraints
    assert len(schema.constraints) == 1
    assert schema.constraints[0].type == "UNIQUENESS"
    assert schema.constraints[0].node_type == "PERSON"
    assert schema.constraints[0].property_names == ("name",)


def test_create_schema_model_invalid_entity(
    schema_builder: SchemaBuilder,
    valid_node_types: Tuple[NodeType, ...],
    valid_relationship_types: Tuple[RelationshipType, ...],
    patterns_with_invalid_entity: tuple[Pattern, ...],
) -> None:
    with pytest.raises(SchemaValidationError) as exc_info:
        schema_builder.create_schema_model(
            list(valid_node_types),
            list(valid_relationship_types),
            list(patterns_with_invalid_entity),
        )
    assert "Node type 'NON_EXISTENT_ENTITY' is not defined" in str(
        exc_info.value
    ), "Should fail due to non-existent entity"


def test_create_schema_model_invalid_relation(
    schema_builder: SchemaBuilder,
    valid_node_types: Tuple[NodeType, ...],
    valid_relationship_types: Tuple[RelationshipType, ...],
    patterns_with_invalid_relation: tuple[Pattern, ...],
) -> None:
    with pytest.raises(SchemaValidationError) as exc_info:
        schema_builder.create_schema_model(
            list(valid_node_types),
            list(valid_relationship_types),
            list(patterns_with_invalid_relation),
        )
    assert "Relationship type 'NON_EXISTENT_RELATION' is not defined" in str(
        exc_info.value
    ), "Should fail due to non-existent relation"


def test_create_schema_model_no_potential_schema(
    schema_builder: SchemaBuilder,
    valid_node_types: Tuple[NodeType, ...],
    valid_relationship_types: Tuple[RelationshipType, ...],
) -> None:
    schema_instance = schema_builder.create_schema_model(
        list(valid_node_types), list(valid_relationship_types)
    )
    assert schema_instance.node_types == valid_node_types
    assert schema_instance.relationship_types == valid_relationship_types
    assert schema_instance.patterns == ()


def test_create_schema_model_no_relations_or_potential_schema(
    schema_builder: SchemaBuilder,
    valid_node_types: Tuple[NodeType, ...],
) -> None:
    schema_instance = schema_builder.create_schema_model(list(valid_node_types))

    assert len(schema_instance.node_types) == 3
    person = schema_instance.node_type_from_label("PERSON")

    assert person is not None
    assert person.description == "An individual human being."
    assert len(person.properties) == 2
    assert person.additional_properties is False

    org = schema_instance.node_type_from_label("ORGANIZATION")
    assert org is not None
    assert org.description == "A structured group of people with a common purpose."
    assert org.additional_properties is True

    age = schema_instance.node_type_from_label("AGE")
    assert age is not None
    assert age.description == "Age of a person in years."
    assert age.additional_properties is True


def test_create_schema_model_missing_relations(
    schema_builder: SchemaBuilder,
    valid_node_types: Tuple[NodeType, ...],
    valid_patterns: tuple[Pattern, ...],
) -> None:
    with pytest.raises(SchemaValidationError) as exc_info:
        schema_builder.create_schema_model(
            node_types=valid_node_types, patterns=valid_patterns
        )
    assert "Relationship types must also be provided when using patterns." in str(
        exc_info.value
    ), "Should fail due to missing relations"


@pytest.fixture
def mock_llm() -> AsyncMock:
    mock = AsyncMock()
    mock.ainvoke = AsyncMock()
    return mock


@pytest.fixture
def valid_schema_json() -> str:
    return """
    {
        "node_types": [
            {
                "label": "Person",
                "properties": [
                    {"name": "name", "type": "STRING"}
                ]
            },
            {
                "label": "Organization",
                "properties": [
                    {"name": "name", "type": "STRING"}
                ]
            }
        ],
        "relationship_types": [
            {
                "label": "WORKS_FOR",
                "properties": [
                    {"name": "since", "type": "DATE"}
                ]
            }
        ],
        "patterns": [
            ["Person", "WORKS_FOR", "Organization"]
        ]
    }
    """


@pytest.fixture
def schema_json_with_valid_constraints() -> str:
    return """
    {
        "node_types": [
            {
                "label": "Person",
                "properties": [
                    {"name": "name", "type": "STRING"},
                    {"name": "email", "type": "STRING"}
                ]
            },
            {
                "label": "Organization",
                "properties": [
                    {"name": "name", "type": "STRING"}
                ]
            }
        ],
        "relationship_types": [
            {
                "label": "WORKS_FOR",
                "properties": [
                    {"name": "since", "type": "DATE"}
                ]
            }
        ],
        "patterns": [
            ["Person", "WORKS_FOR", "Organization"]
        ],
        "constraints": [
            {"type": "UNIQUENESS", "node_type": "Person", "property_name": "name"}
        ]
    }
    """


@pytest.fixture
def schema_json_with_invalid_constraints() -> str:
    return """
    {
        "node_types": [
            {
                "label": "Person",
                "properties": [
                    {"name": "name", "type": "STRING"}
                ]
            },
            {
                "label": "Organization",
                "properties": [
                    {"name": "name", "type": "STRING"}
                ]
            }
        ],
        "relationship_types": [
            {
                "label": "WORKS_FOR",
                "properties": [
                    {"name": "since", "type": "DATE"}
                ]
            }
        ],
        "patterns": [
            ["Person", "WORKS_FOR", "Organization"]
        ],
        "constraints": [
            {"type": "UNIQUENESS", "node_type": "Person", "property_name": "name"},
            {"type": "UNIQUENESS", "node_type": "Person", "property_name": "email"},
            {"type": "UNIQUENESS", "node_type": "NonExistentNode", "property_name": "id"},
            {"type": "UNIQUENESS", "node_type": "Person", "property_name": ""}
        ]
    }
    """


@pytest.fixture
def schema_json_with_null_constraints() -> str:
    return """
    {
        "node_types": [
            {
                "label": "Person",
                "properties": [
                    {"name": "name", "type": "STRING"}
                ]
            },
            {
                "label": "Organization",
                "properties": [
                    {"name": "name", "type": "STRING"}
                ]
            }
        ],
        "relationship_types": [
            {
                "label": "WORKS_FOR",
                "properties": [
                    {"name": "since", "type": "DATE"}
                ]
            }
        ],
        "patterns": [
            ["Person", "WORKS_FOR", "Organization"]
        ],
        "constraints": null
    }
    """


@pytest.fixture
def schema_json_with_required_properties() -> str:
    return """
    {
        "node_types": [
            {
                "label": "Person",
                "properties": [
                    {"name": "name", "type": "STRING", "required": true},
                    {"name": "email", "type": "STRING", "required": false},
                    {"name": "phone", "type": "STRING"}
                ]
            }
        ],
        "relationship_types": [
            {"label": "KNOWS"}
        ],
        "patterns": [
            ["Person", "KNOWS", "Person"]
        ]
    }
    """


@pytest.fixture
def schema_json_with_string_required_values() -> str:
    return """
    {
        "node_types": [
            {
                "label": "Person",
                "properties": [
                    {"name": "name", "type": "STRING", "required": "true"},
                    {"name": "email", "type": "STRING", "required": "yes"},
                    {"name": "phone", "type": "STRING", "required": "false"},
                    {"name": "address", "type": "STRING", "required": "no"}
                ]
            }
        ],
        "relationship_types": [
            {"label": "KNOWS"}
        ],
        "patterns": [
            ["Person", "KNOWS", "Person"]
        ]
    }
    """


@pytest.fixture
def invalid_schema_json() -> str:
    return """
    {
        "node_types": [
            {
                "label": "Person",
            },
        ],
        invalid json content
    }
    """


@pytest.fixture
def schema_from_text(mock_llm: AsyncMock) -> SchemaFromTextExtractor:
    return SchemaFromTextExtractor(llm=mock_llm)


@pytest.mark.asyncio
async def test_schema_from_text_run_valid_response(
    schema_from_text: SchemaFromTextExtractor,
    mock_llm: AsyncMock,
    valid_schema_json: str,
) -> None:
    # configure the mock LLM to return a valid schema JSON
    mock_llm.ainvoke.return_value = LLMResponse(content=valid_schema_json)

    # run the schema extraction
    schema_config = await schema_from_text.run(text="Sample text for extraction")

    # verify the LLM was called with a prompt
    mock_llm.ainvoke.assert_called_once()
    prompt_arg = mock_llm.ainvoke.call_args[0][0]
    assert isinstance(prompt_arg, str)
    assert "Sample text for extraction" in prompt_arg

    # verify the schema was correctly extracted
    assert len(schema_config.node_types) == 2
    assert schema_config.node_type_from_label("Person") is not None
    assert schema_config.node_type_from_label("Organization") is not None

    assert schema_config.relationship_types is not None
    assert schema_config.relationship_type_from_label("WORKS_FOR") is not None

    assert schema_config.patterns is not None
    assert len(schema_config.patterns) == 1
    assert schema_config.patterns[0] == ("Person", "WORKS_FOR", "Organization")


@pytest.mark.asyncio
async def test_schema_from_text_run_invalid_json(
    schema_from_text: SchemaFromTextExtractor,
    mock_llm: AsyncMock,
    invalid_schema_json: str,
) -> None:
    # configure the mock LLM to return invalid JSON
    mock_llm.ainvoke.return_value = LLMResponse(content=invalid_schema_json)

    # verify that running with invalid JSON raises a ValueError
    with pytest.raises(SchemaExtractionError) as exc_info:
        await schema_from_text.run(text="Sample text for extraction")

    assert "not valid JSON" in str(exc_info.value)


@pytest.mark.asyncio
async def test_schema_from_text_custom_template(
    mock_llm: AsyncMock, valid_schema_json: str
) -> None:
    # create a  custom template
    custom_prompt = "This is a custom prompt with text: {text}"
    custom_template = PromptTemplate(template=custom_prompt, expected_inputs=["text"])

    # create SchemaFromTextExtractor with the custom template
    schema_from_text = SchemaFromTextExtractor(
        llm=mock_llm, prompt_template=custom_template
    )

    # configure mock LLM to return valid JSON and capture the prompt that was sent to it
    mock_llm.ainvoke.return_value = LLMResponse(content=valid_schema_json)

    # run the schema extraction
    await schema_from_text.run(text="Sample text")

    # verify the custom prompt was passed to the LLM
    prompt_sent_to_llm = mock_llm.ainvoke.call_args[0][0]
    assert "This is a custom prompt with text" in prompt_sent_to_llm


@pytest.mark.asyncio
async def test_schema_from_text_llm_params(
    mock_llm: AsyncMock, valid_schema_json: str
) -> None:
    # configure custom LLM parameters
    llm_params = {"temperature": 0.1, "max_tokens": 500}

    # create SchemaFromTextExtractor with custom LLM parameters
    schema_from_text = SchemaFromTextExtractor(llm=mock_llm, llm_params=llm_params)

    # configure the mock LLM to return a valid schema JSON
    mock_llm.ainvoke.return_value = LLMResponse(content=valid_schema_json)

    # run the schema extraction
    await schema_from_text.run(text="Sample text")

    # verify the LLM was called with the custom parameters
    mock_llm.ainvoke.assert_called_once()
    call_kwargs = mock_llm.ainvoke.call_args[1]
    assert call_kwargs["temperature"] == 0.1
    assert call_kwargs["max_tokens"] == 500


@pytest.mark.asyncio
async def test_schema_config_save_json(graph_schema: GraphSchema) -> None:
    with tempfile.TemporaryDirectory() as temp_dir:
        # create file path
        json_path = os.path.join(temp_dir, "schema.json")

        # store the schema config
        graph_schema.save(json_path)

        # verify the file exists and has content
        assert os.path.exists(json_path)
        assert os.path.getsize(json_path) > 0

        # verify the content is valid JSON and contains expected data
        with open(json_path, "r") as f:
            data = json.load(f)
            assert "node_types" in data
            assert len(data["node_types"]) == 3


@pytest.mark.asyncio
async def test_schema_config_save_yaml(graph_schema: GraphSchema) -> None:
    with tempfile.TemporaryDirectory() as temp_dir:
        # Create file path
        yaml_path = os.path.join(temp_dir, "schema.yaml")

        # Store the schema config
        graph_schema.save(yaml_path)

        # Verify the file exists and has content
        assert os.path.exists(yaml_path)
        assert os.path.getsize(yaml_path) > 0

        # Verify the content is valid YAML and contains expected data
        with open(yaml_path, "r") as f:
            data = yaml.safe_load(f)
            assert "node_types" in data
            assert len(data["node_types"]) == 3


@pytest.mark.asyncio
async def test_schema_config_from_file(graph_schema: GraphSchema) -> None:
    with tempfile.TemporaryDirectory() as temp_dir:
        # create file paths with different extensions
        json_path = os.path.join(temp_dir, "schema.json")
        yaml_path = os.path.join(temp_dir, "schema.yaml")
        yml_path = os.path.join(temp_dir, "schema.yml")

        # store the schema config in the different formats
        graph_schema.save(json_path)
        graph_schema.save(yaml_path)
        graph_schema.save(yml_path)

        # load using from_file which should detect the format based on extension
        json_schema = GraphSchema.from_file(json_path)
        yaml_schema = GraphSchema.from_file(yaml_path)
        yml_schema = GraphSchema.from_file(yml_path)

        # simple verification that the objects were loaded correctly
        assert isinstance(json_schema, GraphSchema)
        assert isinstance(yaml_schema, GraphSchema)
        assert isinstance(yml_schema, GraphSchema)

        # verify basic structure is intact
        assert "node_types" in json_schema.model_dump()
        assert "node_types" in yaml_schema.model_dump()
        assert "node_types" in yml_schema.model_dump()

        # verify an unsupported extension raises the correct error
        txt_path = os.path.join(temp_dir, "schema.txt")
        graph_schema.save(
            txt_path, format=FileFormat.JSON
        )  # Store as JSON but with .txt extension

        with pytest.raises(ValueError, match="Unsupported file format: None"):
            GraphSchema.from_file(txt_path)


@pytest.fixture
def valid_schema_json_array() -> str:
    return """
    [
        {
            "node_types": [
                {
                    "label": "Person",
                    "properties": [
                        {"name": "name", "type": "STRING"}
                    ]
                },
                {
                    "label": "Organization",
                    "properties": [
                        {"name": "name", "type": "STRING"}
                    ]
                }
            ],
            "relationship_types": [
                {
                    "label": "WORKS_FOR",
                    "properties": [
                        {"name": "since", "type": "DATE"}
                    ]
                }
            ],
            "patterns": [
                ["Person", "WORKS_FOR", "Organization"]
            ]
        }
    ]
    """


@pytest.mark.asyncio
async def test_schema_from_text_run_valid_json_array(
    schema_from_text: SchemaFromTextExtractor,
    mock_llm: AsyncMock,
    valid_schema_json_array: str,
) -> None:
    # configure the mock LLM to return a valid JSON array
    mock_llm.ainvoke.return_value = LLMResponse(content=valid_schema_json_array)

    # run the schema extraction
    schema = await schema_from_text.run(text="Sample text for extraction")

    # verify the schema was correctly extracted from the array
    assert len(schema.node_types) == 2
    assert schema.node_type_from_label("Person") is not None
    assert schema.node_type_from_label("Organization") is not None

    assert schema.relationship_types is not None
    assert schema.relationship_type_from_label("WORKS_FOR") is not None

    assert schema.patterns is not None
    assert len(schema.patterns) == 1
    assert schema.patterns[0] == ("Person", "WORKS_FOR", "Organization")


@pytest.fixture
def schema_json_with_invalid_node_patterns() -> str:
    return """
    {
        "node_types": [
            {
                "label": "Person",
                "properties": [
                    {"name": "name", "type": "STRING"}
                ]
            },
            {
                "label": "Organization",
                "properties": [
                    {"name": "name", "type": "STRING"}
                ]
            }
        ],
        "relationship_types": [
            {
                "label": "WORKS_FOR",
                "properties": [
                    {"name": "since", "type": "DATE"}
                ]
            }
        ],
        "patterns": [
            ["Person", "WORKS_FOR", "Organization"],
            ["Person", "WORKS_FOR", "UndefinedNode"],
            ["UndefinedNode", "WORKS_FOR", "Organization"]
        ]
    }
    """


@pytest.fixture
def schema_json_with_invalid_relationship_patterns() -> str:
    return """
    {
        "node_types": [
            {
                "label": "Person",
                "properties": [
                    {"name": "name", "type": "STRING"}
                ]
            },
            {
                "label": "Organization",
                "properties": [
                    {"name": "name", "type": "STRING"}
                ]
            }
        ],
        "relationship_types": [
            {
                "label": "WORKS_FOR",
                "properties": [
                    {"name": "since", "type": "DATE"}
                ]
            }
        ],
        "patterns": [
            ["Person", "WORKS_FOR", "Organization"],
            ["Person", "UNDEFINED_RELATION", "Organization"],
            ["Organization", "ANOTHER_UNDEFINED_RELATION", "Person"]
        ]
    }
    """


@pytest.fixture
def schema_json_with_nodes_without_labels() -> str:
    return """
    {
        "node_types": [
            {
                "label": "Person",
                "properties": [
                    {"name": "name", "type": "STRING"}
                ]
            },
            {
                "properties": [
                    {"name": "name", "type": "STRING"}
                ]
            },
            {
                "label": "",
                "properties": [
                    {"name": "name", "type": "STRING"}
                ]
            },
            "Organization",
            "",
            "Company",
            "Invalid description with spaces",
            "{\\"invalid\\": \\"json object\\"}"
        ],
        "relationship_types": [
            {
                "label": "WORKS_FOR",
                "properties": [
                    {"name": "since", "type": "DATE"}
                ]
            }
        ],
        "patterns": [
            ["Person", "WORKS_FOR", "Organization"]
        ]
    }
    """


@pytest.fixture
def schema_json_with_relationships_without_labels() -> str:
    return """
    {
        "node_types": [
            {
                "label": "Person",
                "properties": [
                    {"name": "name", "type": "STRING"}
                ]
            },
            {
                "label": "Organization",
                "properties": [
                    {"name": "name", "type": "STRING"}
                ]
            }
        ],
        "relationship_types": [
            {
                "label": "WORKS_FOR",
                "properties": [
                    {"name": "since", "type": "DATE"}
                ]
            },
            {
                "properties": [
                    {"name": "since", "type": "DATE"}
                ]
            },
            {
                "label": "",
                "properties": [
                    {"name": "since", "type": "DATE"}
                ]
            },
            "MANAGES",
            "",
            "SUPERVISES",
            "invalid relationship description",
            "{\\"invalid\\": \\"json\\"}"
        ],
        "patterns": [
            ["Person", "WORKS_FOR", "Organization"],
            ["Person", "MANAGES", "Organization"]
        ]
    }
    """


@pytest.fixture
def schema_json_with_nonexistent_property_constraint() -> str:
    return """
    {
        "node_types": [
            {
                "label": "Person",
                "properties": [
                    {"name": "name", "type": "STRING"}
                ]
            }
        ],
        "relationship_types": [],
        "patterns": [],
        "constraints": [
            {"type": "UNIQUENESS", "node_type": "Person", "property_name": "name"},
            {"type": "UNIQUENESS", "node_type": "Person", "property_name": "nonexistent_property"}
        ]
    }
    """


@pytest.mark.asyncio
async def test_schema_from_text_filters_invalid_node_patterns(
    schema_from_text: SchemaFromTextExtractor,
    mock_llm: AsyncMock,
    schema_json_with_invalid_node_patterns: str,
) -> None:
    # configure the mock LLM to return schema with invalid node patterns
    mock_llm.ainvoke.return_value = LLMResponse(
        content=schema_json_with_invalid_node_patterns
    )

    # run the schema extraction
    schema = await schema_from_text.run(text="Sample text for extraction")

    # verify that invalid node patterns were filtered out (2 out of 3 patterns should be removed)
    assert schema.patterns is not None
    assert len(schema.patterns) == 1
    assert schema.patterns[0] == ("Person", "WORKS_FOR", "Organization")


@pytest.mark.asyncio
async def test_schema_from_text_filters_invalid_relationship_patterns(
    schema_from_text: SchemaFromTextExtractor,
    mock_llm: AsyncMock,
    schema_json_with_invalid_relationship_patterns: str,
) -> None:
    # configure the mock LLM to return schema with invalid relationship patterns
    mock_llm.ainvoke.return_value = LLMResponse(
        content=schema_json_with_invalid_relationship_patterns
    )

    # run the schema extraction
    schema = await schema_from_text.run(text="Sample text for extraction")

    # verify that invalid relationship patterns were filtered out (2 out of 3 patterns should be removed)
    assert schema.patterns is not None
    assert len(schema.patterns) == 1
    assert schema.patterns[0] == ("Person", "WORKS_FOR", "Organization")


@pytest.mark.asyncio
async def test_schema_from_text_filters_nodes_without_labels(
    schema_from_text: SchemaFromTextExtractor,
    mock_llm: AsyncMock,
    schema_json_with_nodes_without_labels: str,
) -> None:
    # configure the mock LLM to return schema with nodes without labels
    mock_llm.ainvoke.return_value = LLMResponse(
        content=schema_json_with_nodes_without_labels
    )

    # run the schema extraction
    schema = await schema_from_text.run(text="Sample text for extraction")

    # verify that nodes without labels were filtered out (5 out of 8 nodes should be removed)
    assert len(schema.node_types) == 3
    assert schema.node_type_from_label("Person") is not None
    assert schema.node_type_from_label("Organization") is not None
    assert schema.node_type_from_label("Company") is not None

    # verify that the pattern is still valid with the remaining nodes
    assert schema.patterns is not None
    assert len(schema.patterns) == 1
    assert schema.patterns[0] == ("Person", "WORKS_FOR", "Organization")


@pytest.mark.asyncio
async def test_schema_from_text_filters_relationships_without_labels(
    schema_from_text: SchemaFromTextExtractor,
    mock_llm: AsyncMock,
    schema_json_with_relationships_without_labels: str,
) -> None:
    # configure the mock LLM to return schema with relationships without labels
    mock_llm.ainvoke.return_value = LLMResponse(
        content=schema_json_with_relationships_without_labels
    )

    # run the schema extraction
    schema = await schema_from_text.run(text="Sample text for extraction")

    # verify that relationships without labels were filtered out (5 out of 8 relationships should be removed)
    assert schema.relationship_types is not None
    assert len(schema.relationship_types) == 3
    assert schema.relationship_type_from_label("WORKS_FOR") is not None
    assert schema.relationship_type_from_label("MANAGES") is not None
    assert schema.relationship_type_from_label("SUPERVISES") is not None

    # verify that the patterns are still valid with the remaining relationships
    assert schema.patterns is not None
    assert len(schema.patterns) == 2
    assert ("Person", "WORKS_FOR", "Organization") in schema.patterns
    assert ("Person", "MANAGES", "Organization") in schema.patterns


@pytest.mark.asyncio
async def test_schema_from_text_with_valid_constraints(
    schema_from_text: SchemaFromTextExtractor,
    mock_llm: AsyncMock,
    schema_json_with_valid_constraints: str,
) -> None:
    # configure the mock LLM to return schema with valid constraints
    mock_llm.ainvoke.return_value = LLMResponse(
        content=schema_json_with_valid_constraints
    )

    # run the schema extraction
    schema = await schema_from_text.run(text="Sample text for extraction")

    assert len(schema.constraints) == 1
    assert schema.constraints[0].type == "UNIQUENESS"
    assert schema.constraints[0].node_type == "Person"
    assert schema.constraints[0].property_names == ("name",)


@pytest.mark.asyncio
async def test_schema_from_text_filters_invalid_constraints(
    schema_from_text: SchemaFromTextExtractor,
    mock_llm: AsyncMock,
    schema_json_with_invalid_constraints: str,
) -> None:
    # configure the mock LLM to return schema with invalid constraints
    mock_llm.ainvoke.return_value = LLMResponse(
        content=schema_json_with_invalid_constraints
    )

    # run the schema extraction
    schema = await schema_from_text.run(text="Sample text for extraction")

    # verify that invalid constraints were filtered out:
    # constraints with NonExistentNode should be removed
    # constraint with empty property_name should be removed
    # only the valid constraint should remain
    assert len(schema.constraints) == 1
    assert schema.constraints[0].node_type == "Person"
    assert schema.constraints[0].property_names == ("name",)


@pytest.mark.asyncio
async def test_schema_from_text_filters_constraint_with_nonexistent_property(
    schema_from_text: SchemaFromTextExtractor,
    mock_llm: AsyncMock,
    schema_json_with_nonexistent_property_constraint: str,
) -> None:
    # configure the mock LLM to return schema with constraint on nonexistent property
    mock_llm.ainvoke.return_value = LLMResponse(
        content=schema_json_with_nonexistent_property_constraint
    )

    # run the schema extraction
    schema = await schema_from_text.run(text="Sample text for extraction")

    # verify that only the valid constraint (with "name" property) remains
    # the constraint with "nonexistent_property" should be filtered out
    assert len(schema.constraints) == 1
    assert schema.constraints[0].property_names == ("name",)


@pytest.mark.asyncio
async def test_schema_from_text_handles_null_constraints(
    schema_from_text: SchemaFromTextExtractor,
    mock_llm: AsyncMock,
    schema_json_with_null_constraints: str,
) -> None:
    # configure the mock LLM to return schema with null constraints
    mock_llm.ainvoke.return_value = LLMResponse(
        content=schema_json_with_null_constraints
    )

    # run the schema extraction - should not crash
    schema = await schema_from_text.run(text="Sample text for extraction")

    # verify schema was created with empty constraints
    assert len(schema.constraints) == 0


@pytest.mark.asyncio
async def test_schema_from_text_handles_missing_constraints(
    schema_from_text: SchemaFromTextExtractor,
    mock_llm: AsyncMock,
    valid_schema_json: str,
) -> None:
    # configure the mock LLM to return schema without constraints field
    mock_llm.ainvoke.return_value = LLMResponse(content=valid_schema_json)

    #  run the schema extraction - should not crash
    schema = await schema_from_text.run(text="Sample text for extraction")

    # verify schema was created with empty constraints
    assert len(schema.constraints) == 0


def test_clean_json_content_markdown_with_json_language(
    schema_from_text: SchemaFromTextExtractor,
) -> None:
    content = """```json
{"node_types": [{"label": "Person"}]}
```"""

    cleaned = schema_from_text._clean_json_content(content)
    assert cleaned == '{"node_types": [{"label": "Person"}]}'


def test_clean_json_content_markdown_without_language(
    schema_from_text: SchemaFromTextExtractor,
) -> None:
    content = """```
{"node_types": [{"label": "Person"}]}
```"""

    cleaned = schema_from_text._clean_json_content(content)
    assert cleaned == '{"node_types": [{"label": "Person"}]}'


def test_clean_json_content_plain_json(
    schema_from_text: SchemaFromTextExtractor,
) -> None:
    content = '{"node_types": [{"label": "Person"}]}'

    cleaned = schema_from_text._clean_json_content(content)
    assert cleaned == '{"node_types": [{"label": "Person"}]}'


@pytest.mark.asyncio
async def test_schema_from_text_with_required_properties(
    schema_from_text: SchemaFromTextExtractor,
    mock_llm: AsyncMock,
    schema_json_with_required_properties: str,
) -> None:
    mock_llm.ainvoke.return_value = LLMResponse(
        content=schema_json_with_required_properties
    )

    schema = await schema_from_text.run(text="Sample text for test")

    person = schema.node_type_from_label("Person")
    assert person is not None

    # Legacy required:true migrates to EXISTENCE; flags are cleared on PropertyType
    assert schema.existence_property_names_for_node("Person") == {"name"}
    name_prop = next((p for p in person.properties if p.name == "name"), None)
    email_prop = next((p for p in person.properties if p.name == "email"), None)
    phone_prop = next((p for p in person.properties if p.name == "phone"), None)

    assert name_prop is not None and name_prop.required is False
    assert email_prop is not None and email_prop.required is False
    assert phone_prop is not None and phone_prop.required is False


@pytest.mark.asyncio
async def test_schema_from_text_string_required_coerced_without_existence_migration(
    schema_from_text: SchemaFromTextExtractor,
    mock_llm: AsyncMock,
    schema_json_with_string_required_values: str,
) -> None:
    """LLMs may emit string truthiness for ``required``; Pydantic coerces it on ``PropertyType``.

    Migration of legacy ``required`` to ``EXISTENCE`` constraints only runs when the raw
    property dict has ``required is True`` (JSON boolean). String values such as ``\"true\"``
    are not migrated; add ``ConstraintType`` rows with type ``EXISTENCE`` on ``GraphSchema``
    if you need existence semantics for that case.
    """
    mock_llm.ainvoke.return_value = LLMResponse(
        content=schema_json_with_string_required_values
    )

    schema = await schema_from_text.run(text="Sample text for test")

    person = schema.node_type_from_label("Person")
    assert person is not None

    assert schema.existence_property_names_for_node("Person") == set()
    name_prop = next((p for p in person.properties if p.name == "name"), None)
    email_prop = next((p for p in person.properties if p.name == "email"), None)
    assert name_prop is not None
    assert email_prop is not None
    assert name_prop.model_dump().get("required") is True
    assert email_prop.model_dump().get("required") is True

    phone_prop = next((p for p in person.properties if p.name == "phone"), None)
    address_prop = next((p for p in person.properties if p.name == "address"), None)
    assert phone_prop is not None
    assert address_prop is not None
    assert phone_prop.model_dump().get("required") is False
    assert address_prop.model_dump().get("required") is False


@pytest.mark.asyncio
async def test_schema_from_text_handles_missing_required_field(
    schema_from_text: SchemaFromTextExtractor,
    mock_llm: AsyncMock,
    valid_schema_json: str,
) -> None:
    mock_llm.ainvoke.return_value = LLMResponse(content=valid_schema_json)

    schema = await schema_from_text.run(text="Sample text")

    person = schema.node_type_from_label("Person")
    assert person is not None

    # All properties should have required=False (default)
    for prop in person.properties:
        assert prop.required is False


@pytest.mark.asyncio
async def test_schema_from_text_uniqueness_does_not_force_required_property(
    schema_from_text: SchemaFromTextExtractor,
    mock_llm: AsyncMock,
) -> None:
    schema_json = """
    {
        "node_types": [
            {
                "label": "Person",
                "properties": [
                    {"name": "name", "type": "STRING", "required": false},
                    {"name": "email", "type": "STRING", "required": false}
                ]
            }
        ],
        "relationship_types": [],
        "patterns": [],
        "constraints": [
            {"type": "UNIQUENESS", "node_type": "Person", "property_name": "name"}
        ]
    }
    """
    mock_llm.ainvoke.return_value = LLMResponse(content=schema_json)

    schema = await schema_from_text.run(text="Sample text")

    person = schema.node_type_from_label("Person")
    assert person is not None

    name_prop = next((p for p in person.properties if p.name == "name"), None)
    email_prop = next((p for p in person.properties if p.name == "email"), None)

    assert name_prop is not None and name_prop.required is False
    assert email_prop is not None and email_prop.required is False
    assert schema.existence_property_names_for_node("Person") == set()
    assert schema.constraints[0].type == "UNIQUENESS"


@pytest.mark.asyncio
@patch("neo4j_graphrag.experimental.components.schema.get_structured_schema")
async def test_schema_from_existing_graph(mock_get_structured_schema: Mock) -> None:
    mock_get_structured_schema.return_value = {
        "node_props": {
            "Person": [
                {"property": "id", "type": "INTEGER"},
                {"property": "name", "type": "STRING"},
            ]
        },
        "rel_props": {"KNOWS": [{"property": "fromDate", "type": "DATE"}]},
        "relationships": [
            {"start": "Person", "type": "KNOWS", "end": "Person"},
            {"start": "Person", "type": "LIVES_IN", "end": "City"},
        ],
        "metadata": {
            "constraint": [
                {
                    "id": 7,
                    "name": "person_id",
                    "type": "NODE_PROPERTY_EXISTENCE",
                    "entityType": "NODE",
                    "labelsOrTypes": ["Person"],
                    "properties": ["id"],
                    "ownedIndex": "person_id",
                    "propertyType": None,
                },
            ],
            "index": [
                {
                    "label": "Person",
                    "properties": ["name"],
                    "size": 2,
                    "type": "RANGE",
                    "valuesSelectivity": 1.0,
                    "distinctValues": 2.0,
                },
            ],
        },
    }
    driver = Mock()
    schema_builder = SchemaFromExistingGraphExtractor(
        driver=driver,
    )
    schema = await schema_builder.run()
    assert isinstance(schema, GraphSchema)
    assert len(schema.node_types) == 2
    person_node_type = schema.node_type_from_label("Person")
    assert person_node_type is not None
    id_person_property = [p for p in person_node_type.properties if p.name == "id"][0]
    assert id_person_property.required is False
    assert schema.existence_property_names_for_node("Person") == {"id"}
    assert any(
        c.type == "EXISTENCE"
        and c.node_type == "Person"
        and c.property_names == ("id",)
        for c in schema.constraints
    )
    assert person_node_type.additional_properties is False
    city_node_type = schema.node_type_from_label("City")
    assert city_node_type is not None
    assert city_node_type.additional_properties is True

    assert len(schema.relationship_types) == 2
    knows_rel = schema.relationship_type_from_label("KNOWS")
    assert knows_rel is not None
    assert knows_rel.additional_properties is False
    lives_rel = schema.relationship_type_from_label("LIVES_IN")
    assert lives_rel is not None
    assert lives_rel.additional_properties is True

    assert schema.additional_node_types is False
    assert schema.additional_relationship_types is False
    assert schema.additional_patterns is False

    assert schema.patterns == (
        Pattern(source="Person", relationship="KNOWS", target="Person"),
        Pattern(source="Person", relationship="LIVES_IN", target="City"),
    )


@pytest.mark.asyncio
@patch("neo4j_graphrag.experimental.components.schema.get_structured_schema")
async def test_schema_from_existing_graph_additional_params(
    mock_get_structured_schema: Mock,
) -> None:
    mock_get_structured_schema.return_value = {
        "node_props": {
            "Person": [
                {"property": "id", "type": "INTEGER"},
                {"property": "name", "type": "STRING"},
            ]
        },
        "rel_props": {"KNOWS": [{"property": "fromDate", "type": "DATE"}]},
        "relationships": [
            {"start": "Person", "type": "KNOWS", "end": "Person"},
            {"start": "Person", "type": "LIVES_IN", "end": "City"},
        ],
    }
    driver = Mock()
    schema_builder = SchemaFromExistingGraphExtractor(
        driver=driver,
        additional_node_types=True,
        additional_relationship_types=True,
        additional_patterns=True,
        additional_properties=True,
    )
    schema = await schema_builder.run()
    person_node_type = schema.node_type_from_label("Person")
    assert person_node_type is not None
    assert person_node_type.additional_properties is True
    city_node_type = schema.node_type_from_label("City")
    assert city_node_type is not None
    assert city_node_type.additional_properties is True

    knows_rel = schema.relationship_type_from_label("KNOWS")
    assert knows_rel is not None
    assert knows_rel.additional_properties is True
    lives_rel = schema.relationship_type_from_label("LIVES_IN")
    assert lives_rel is not None
    assert lives_rel.additional_properties is True

    assert schema.additional_node_types is True
    assert schema.additional_relationship_types is True
    assert schema.additional_patterns is True


def test_extract_graph_constraints_from_metadata_node_key_maps_to_key() -> None:
    """Neo4j ``NODE_KEY`` metadata must become ``GraphConstraintType.KEY``, not EXISTENCE."""
    structured_schema: dict[str, Any] = {
        "metadata": {
            "constraint": [
                {
                    "type": "NODE_KEY",
                    "labelsOrTypes": ["Person"],
                    "properties": ["email"],
                }
            ]
        }
    }
    out = SchemaFromExistingGraphExtractor._extract_graph_constraints_from_metadata(
        structured_schema
    )
    assert out == [
        {
            "type": GraphConstraintType.KEY.value,
            "node_type": "Person",
            "property_name": "email",
            "property_names": ("email",),
            "relationship_type": None,
        }
    ]


def test_extract_graph_constraints_from_metadata_relationship_key_maps_to_key() -> None:
    """Neo4j ``RELATIONSHIP_KEY`` metadata maps to relationship-scoped ``KEY`` constraints."""
    structured_schema: dict[str, Any] = {
        "metadata": {
            "constraint": [
                {
                    "type": "RELATIONSHIP_KEY",
                    "labelsOrTypes": ["WORKS_FOR"],
                    "properties": ["since"],
                }
            ]
        }
    }
    out = SchemaFromExistingGraphExtractor._extract_graph_constraints_from_metadata(
        structured_schema
    )
    assert out == [
        {
            "type": GraphConstraintType.KEY.value,
            "node_type": "",
            "property_name": "since",
            "property_names": ("since",),
            "relationship_type": "WORKS_FOR",
        }
    ]


def test_graph_schema_extraction_output_json_schema_lean_root() -> None:
    """Structured-output schema must not include pipeline-only GraphSchema flags."""
    from neo4j_graphrag.experimental.components.graph_schema_extraction import (
        GraphSchemaExtractionOutput,
    )

    raw = GraphSchemaExtractionOutput.model_json_schema()
    dumped = json.dumps(raw)
    assert "additional_node_types" not in dumped
    assert "additional_relationship_types" not in dumped
    assert "additional_patterns" not in dumped


def test_graph_schema_extraction_constraint_schema_avoids_null_type_for_vertex() -> (
    None
):
    """Vertex AI maps JSON Schema to protobuf and rejects ``type: \"null\"`` (e.g. ``Optional``)."""
    from neo4j_graphrag.experimental.components.graph_schema_extraction import (
        GraphSchemaExtractionOutput,
    )

    raw = GraphSchemaExtractionOutput.model_json_schema()
    ect = (raw.get("$defs") or {}).get("ExtractedConstraintType")
    assert ect is not None
    rel_schema = (ect.get("properties") or {}).get("relationship_type", {})
    assert "anyOf" not in rel_schema
    assert '"type": "null"' not in json.dumps(rel_schema)


def test_graph_schema_from_extraction_output() -> None:
    from neo4j_graphrag.experimental.components.graph_schema_extraction import (
        ExtractedConstraintType,
        ExtractedNodeType,
        ExtractedPropertyType,
        GraphSchemaExtractionOutput,
    )

    dto = GraphSchemaExtractionOutput(
        node_types=[
            ExtractedNodeType(
                label="Person",
                properties=[
                    ExtractedPropertyType(name="name", type="STRING"),
                ],
            )
        ],
        relationship_types=[],
        patterns=[],
        constraints=[
            ExtractedConstraintType(
                type="UNIQUENESS",
                node_type="Person",
                property_names=["name"],
            ),
            ExtractedConstraintType(
                type="EXISTENCE",
                node_type="Person",
                property_names=["name"],
                relationship_type="",
            ),
        ],
    )
    gs = GraphSchema.from_extraction_output(dto)
    assert gs.node_types[0].label == "Person"
    assert gs.node_types[0].properties[0].required is False
    assert gs.existence_property_names_for_node("Person") == {"name"}
    assert {c.property_names for c in gs.constraints} == {("name",)}
    assert gs.additional_node_types is False


def test_validate_extraction_dict_to_graph_schema() -> None:
    from neo4j_graphrag.experimental.components.schema import (
        validate_extraction_dict_to_graph_schema,
    )

    d = {
        "node_types": [
            {
                "label": "Person",
                "properties": [{"name": "name", "type": "STRING", "required": False}],
            }
        ],
        "relationship_types": [],
        "patterns": [],
        "constraints": [],
    }
    gs = validate_extraction_dict_to_graph_schema(d)
    assert len(gs.node_types) == 1
    assert gs.node_types[0].label == "Person"