"""
Tests for JSON schema constraint functionality used by JsonArrayParser
"""

import json
import unittest

import jsonschema

from sglang.srt.entrypoints.openai.protocol import (
    Function,
    Tool,
    ToolChoice,
    ToolChoiceFuncName,
)
from sglang.srt.function_call.function_call_parser import FunctionCallParser
from sglang.srt.function_call.utils import (
    _get_tool_schema_defs,
    get_json_schema_constraint,
)


class TestJsonSchemaConstraint(unittest.TestCase):
    """Test JSON schema constraint generation for tool choices"""

    def setUp(self):
        """Set up test tools"""
        self.tools = [
            Tool(
                type="function",
                function=Function(
                    name="get_weather",
                    description="Get weather information",
                    parameters={
                        "type": "object",
                        "properties": {
                            "location": {
                                "type": "string",
                                "description": "Location to get weather for",
                            },
                            "unit": {
                                "type": "string",
                                "enum": ["celsius", "fahrenheit"],
                                "description": "Temperature unit",
                            },
                        },
                        "required": ["location"],
                    },
                ),
            ),
            Tool(
                type="function",
                function=Function(
                    name="search",
                    description="Search for information",
                    parameters={
                        "type": "object",
                        "properties": {
                            "query": {
                                "type": "string",
                                "description": "Search query",
                            },
                        },
                        "required": ["query"],
                    },
                ),
            ),
        ]

    def test_required_tool_choice_schema(self):
        """Test schema generation for tool_choice='required'"""
        schema = get_json_schema_constraint(self.tools, "required")

        self.assertIsNotNone(schema)
        jsonschema.Draft202012Validator.check_schema(schema)

        self.assertEqual(schema["type"], "array")
        self.assertEqual(schema["minItems"], 1)
        self.assertIn("items", schema)
        self.assertIn("anyOf", schema["items"])

        # Should have schemas for both tools
        self.assertEqual(len(schema["items"]["anyOf"]), 2)

        # Check that each tool schema is present
        tool_names = [
            item["properties"]["name"]["enum"][0] for item in schema["items"]["anyOf"]
        ]
        self.assertIn("get_weather", tool_names)
        self.assertIn("search", tool_names)

    def test_specific_tool_choice_schema(self):
        """Test schema generation for specific tool choice"""
        tool_choice = ToolChoice(
            type="function", function=ToolChoiceFuncName(name="get_weather")
        )
        schema = get_json_schema_constraint(self.tools, tool_choice)

        self.assertIsNotNone(schema)
        jsonschema.Draft202012Validator.check_schema(schema)

        self.assertEqual(schema["type"], "array")
        self.assertEqual(schema["minItems"], 1)
        self.assertEqual(schema["maxItems"], 1)

        # Should only have schema for the specific tool
        item_schema = schema["items"]
        self.assertEqual(item_schema["properties"]["name"]["enum"], ["get_weather"])
        self.assertIn("parameters", item_schema["properties"])

    def test_specific_tool_choice_dict_schema(self):
        """Test schema generation for specific tool choice as ToolChoice object"""
        tool_choice = ToolChoice(
            type="function", function=ToolChoiceFuncName(name="search")
        )
        schema = get_json_schema_constraint(self.tools, tool_choice)

        self.assertIsNotNone(schema)
        jsonschema.Draft202012Validator.check_schema(schema)

        self.assertEqual(schema["type"], "array")
        self.assertEqual(schema["minItems"], 1)
        self.assertEqual(schema["maxItems"], 1)

        # Should only have schema for the specific tool
        item_schema = schema["items"]
        self.assertEqual(item_schema["properties"]["name"]["enum"], ["search"])
        self.assertIn("parameters", item_schema["properties"])

    def test_nonexistent_tool_choice(self):
        """Test schema generation for nonexistent tool"""
        tool_choice = ToolChoice(
            type="function", function=ToolChoiceFuncName(name="nonexistent")
        )
        schema = get_json_schema_constraint(self.tools, tool_choice)

        self.assertIsNone(schema)

    def test_nonexistent_tool_choice_dict(self):
        """Test schema generation for nonexistent tool as dict"""
        tool_choice = {"type": "function", "function": {"name": "nonexistent"}}
        schema = get_json_schema_constraint(self.tools, tool_choice)

        self.assertIsNone(schema)

    def test_auto_tool_choice_schema(self):
        """Test schema generation for tool_choice='auto'"""
        schema = get_json_schema_constraint(self.tools, "auto")

        self.assertIsNone(schema)

    def test_none_tool_choice_schema(self):
        """Test schema generation for tool_choice=None"""
        schema = get_json_schema_constraint(self.tools, None)

        self.assertIsNone(schema)

    def test_tools_with_defs(self):
        """Test schema generation with tools that have $defs"""
        tools_with_defs = [
            Tool(
                type="function",
                function=Function(
                    name="complex_tool",
                    description="Tool with complex schema",
                    parameters={
                        "type": "object",
                        "properties": {
                            "data": {
                                "type": "object",
                                "properties": {
                                    "nested": {"$ref": "#/$defs/NestedType"},
                                },
                            },
                        },
                        "$defs": {
                            "NestedType": {
                                "type": "object",
                                "properties": {
                                    "value": {"type": "string"},
                                },
                            },
                        },
                    },
                ),
            ),
        ]

        try:
            _get_tool_schema_defs(tools_with_defs)
        except ValueError as e:
            self.fail(f"Should not raise ValueError, but got: {e}")

        schema = get_json_schema_constraint(tools_with_defs, "required")

        self.assertIsNotNone(schema)
        jsonschema.Draft202012Validator.check_schema(schema)

        self.assertIn("$defs", schema)
        self.assertIn("NestedType", schema["$defs"])

    def test_tools_without_parameters(self):
        """Test schema generation with tools that have no parameters"""
        tools_without_params = [
            Tool(
                type="function",
                function=Function(
                    name="simple_tool",
                    description="Tool without parameters",
                    parameters=None,
                ),
            ),
        ]

        schema = get_json_schema_constraint(tools_without_params, "required")

        self.assertIsNotNone(schema)
        jsonschema.Draft202012Validator.check_schema(schema)

        item_schema = schema["items"]["anyOf"][0]
        self.assertEqual(
            item_schema["properties"]["parameters"],
            {"type": "object", "properties": {}},
        )

    def test_json_schema_vs_ebnf_constraint_generation(self):
        """Test direct comparison between JSON schema and EBNF constraint generation"""

        # Test with specific tool choice
        tool_choice = ToolChoice(
            type="function", function=ToolChoiceFuncName(name="get_weather")
        )

        # Generate JSON schema constraint
        json_schema = get_json_schema_constraint(self.tools, tool_choice)

        self.assertIsNotNone(json_schema)
        jsonschema.Draft202012Validator.check_schema(json_schema)

        # Generate EBNF constraint using FunctionCallParser
        parser = FunctionCallParser(
            self.tools, "llama3"
        )  # Use a parser that supports EBNF
        ebnf_constraint = parser.get_ebnf(tool_choice)

        # Verify JSON schema constraint
        self.assertEqual(json_schema["type"], "array")
        self.assertEqual(json_schema["minItems"], 1)
        self.assertEqual(json_schema["maxItems"], 1)

        # Verify EBNF constraint
        self.assertIsNotNone(ebnf_constraint)
        self.assertIsInstance(ebnf_constraint, str)
        self.assertIn("get_weather", ebnf_constraint)

        # Test with required tool choice
        required_json_schema = get_json_schema_constraint(self.tools, "required")

        self.assertIsNotNone(required_json_schema)
        jsonschema.Draft202012Validator.check_schema(required_json_schema)

        required_ebnf_constraint = parser.get_ebnf("required")

        # Verify required JSON schema constraint
        self.assertEqual(required_json_schema["type"], "array")
        self.assertEqual(required_json_schema["minItems"], 1)
        self.assertIn("anyOf", required_json_schema["items"])

        # Verify required EBNF constraint
        self.assertIsNotNone(required_ebnf_constraint)
        self.assertIsInstance(required_ebnf_constraint, str)

        # Both should contain references to the available tools
        tool_names = [tool.function.name for tool in self.tools]
        for tool_name in tool_names:
            self.assertIn(tool_name, required_ebnf_constraint)

    def test_conflicting_defs_raises_valueerror(self):
        """Test that conflicting tool definitions raise ValueError with proper message"""
        tools_with_conflicting_defs = [
            Tool(
                type="function",
                function=Function(
                    name="tool1",
                    description="Tool 1",
                    parameters={
                        "type": "object",
                        "properties": {},
                        "$defs": {
                            "ConflictingType": {
                                "type": "object",
                                "properties": {"value": {"type": "string"}},
                            },
                        },
                    },
                ),
            ),
            Tool(
                type="function",
                function=Function(
                    name="tool2",
                    description="Tool 2",
                    parameters={
                        "type": "object",
                        "properties": {},
                        "$defs": {
                            "ConflictingType": {
                                "type": "object",
                                "properties": {"value": {"type": "number"}},
                            },
                        },
                    },
                ),
            ),
        ]

        with self.assertRaises(ValueError) as context:
            _get_tool_schema_defs(tools_with_conflicting_defs)

        self.assertIn(
            "Tool definition 'ConflictingType' has multiple schemas",
            str(context.exception),
        )
        self.assertIn("which is not supported", str(context.exception))

    def test_tools_with_empty_defs(self):
        """Test tools with empty $defs objects"""
        tools_with_empty_defs = [
            Tool(
                type="function",
                function=Function(
                    name="empty_defs_tool",
                    description="Tool with empty $defs",
                    parameters={
                        "type": "object",
                        "properties": {
                            "data": {"type": "string"},
                        },
                        "required": ["data"],
                        "$defs": {},
                    },
                ),
            ),
        ]

        try:
            _get_tool_schema_defs(tools_with_empty_defs)
        except ValueError as e:
            self.fail(f"Should not raise ValueError, but got: {e}")

        schema = get_json_schema_constraint(tools_with_empty_defs, "required")
        self.assertIsNotNone(schema)
        jsonschema.Draft202012Validator.check_schema(schema)

        # Should not have $defs section when empty
        self.assertNotIn("$defs", schema)

    def test_tools_with_identical_defs(self):
        """Test different tools with same $defs names but identical schemas (should not raise exception)"""
        tools_with_identical_defs = [
            Tool(
                type="function",
                function=Function(
                    name="weather_tool",
                    description="Get weather information",
                    parameters={
                        "type": "object",
                        "properties": {
                            "location": {"$ref": "#/$defs/Location"},
                        },
                        "required": ["location"],
                        "$defs": {
                            "Location": {
                                "type": "object",
                                "properties": {
                                    "lat": {"type": "number"},
                                    "lon": {"type": "number"},
                                },
                                "required": ["lat", "lon"],
                            },
                        },
                    },
                ),
            ),
            Tool(
                type="function",
                function=Function(
                    name="address_tool",
                    description="Get address information",
                    parameters={
                        "type": "object",
                        "properties": {
                            "address": {"$ref": "#/$defs/Location"},
                        },
                        "required": ["address"],
                        "$defs": {
                            "Location": {
                                "type": "object",
                                "properties": {
                                    "lat": {"type": "number"},
                                    "lon": {"type": "number"},
                                },
                                "required": ["lat", "lon"],
                            },
                        },
                    },
                ),
            ),
        ]

        try:
            _get_tool_schema_defs(tools_with_identical_defs)
        except ValueError as e:
            self.fail(
                f"Should not raise ValueError for identical schemas, but got: {e}"
            )

        # Also test that schema generation works
        schema = get_json_schema_constraint(tools_with_identical_defs, "required")
        self.assertIsNotNone(schema)
        jsonschema.Draft202012Validator.check_schema(schema)

        # Verify both tools are present
        tool_names = [
            item["properties"]["name"]["enum"][0] for item in schema["items"]["anyOf"]
        ]
        self.assertIn("weather_tool", tool_names)
        self.assertIn("address_tool", tool_names)

        # Should have $defs with Location
        self.assertIn("$defs", schema)
        self.assertIn("Location", schema["$defs"])

    def test_tools_with_nested_defs(self):
        """Test tools with nested $defs"""
        tools_with_nested_defs = [
            Tool(
                type="function",
                function=Function(
                    name="complex_tool",
                    description="Tool with nested $defs",
                    parameters={
                        "type": "object",
                        "properties": {
                            "user": {"$ref": "#/$defs/User"},
                            "settings": {"$ref": "#/$defs/Settings"},
                        },
                        "required": ["user"],
                        "$defs": {
                            "User": {
                                "type": "object",
                                "properties": {
                                    "id": {"type": "string"},
                                    "profile": {"$ref": "#/$defs/Profile"},
                                },
                                "required": ["id"],
                            },
                            "Profile": {
                                "type": "object",
                                "properties": {
                                    "name": {"type": "string"},
                                    "email": {"type": "string", "format": "email"},
                                },
                                "required": ["name"],
                            },
                            "Settings": {
                                "type": "object",
                                "properties": {
                                    "theme": {
                                        "type": "string",
                                        "enum": ["light", "dark"],
                                    },
                                    "notifications": {"type": "boolean"},
                                },
                            },
                        },
                    },
                ),
            ),
        ]

        try:
            _get_tool_schema_defs(tools_with_nested_defs)
        except ValueError as e:
            self.fail(f"Should not raise ValueError, but got: {e}")

        schema = get_json_schema_constraint(tools_with_nested_defs, "required")
        self.assertIsNotNone(schema)
        jsonschema.Draft202012Validator.check_schema(schema)

        # Verify all $defs are properly included
        self.assertIn("$defs", schema)
        self.assertIn("User", schema["$defs"])
        self.assertIn("Profile", schema["$defs"])
        self.assertIn("Settings", schema["$defs"])

    def test_mixed_tools_with_and_without_defs(self):
        """Test mixed tools with and without $defs"""
        mixed_tools = [
            Tool(
                type="function",
                function=Function(
                    name="simple_tool",
                    description="Simple tool without $defs",
                    parameters={
                        "type": "object",
                        "properties": {
                            "query": {"type": "string"},
                        },
                        "required": ["query"],
                    },
                ),
            ),
            Tool(
                type="function",
                function=Function(
                    name="complex_tool",
                    description="Complex tool with $defs",
                    parameters={
                        "type": "object",
                        "properties": {
                            "data": {"$ref": "#/$defs/DataType"},
                        },
                        "required": ["data"],
                        "$defs": {
                            "DataType": {
                                "type": "object",
                                "properties": {
                                    "value": {"type": "string"},
                                    "metadata": {"type": "object"},
                                },
                                "required": ["value"],
                            },
                        },
                    },
                ),
            ),
            Tool(
                type="function",
                function=Function(
                    name="another_simple_tool",
                    description="Another simple tool",
                    parameters={
                        "type": "object",
                        "properties": {
                            "id": {"type": "integer"},
                        },
                        "required": ["id"],
                    },
                ),
            ),
        ]

        try:
            _get_tool_schema_defs(mixed_tools)
        except ValueError as e:
            self.fail(f"Should not raise ValueError, but got: {e}")

        schema = get_json_schema_constraint(mixed_tools, "required")
        self.assertIsNotNone(schema)
        jsonschema.Draft202012Validator.check_schema(schema)

        # Should have $defs from the complex tool
        self.assertIn("$defs", schema)
        self.assertIn("DataType", schema["$defs"])

        # Should have all three tools
        tool_names = [
            item["properties"]["name"]["enum"][0] for item in schema["items"]["anyOf"]
        ]
        self.assertEqual(len(tool_names), 3)
        self.assertIn("simple_tool", tool_names)
        self.assertIn("complex_tool", tool_names)
        self.assertIn("another_simple_tool", tool_names)

    def test_tools_with_defs_but_no_refs(self):
        """Test tools with $defs but no $ref usage"""
        tools_with_unused_defs = [
            Tool(
                type="function",
                function=Function(
                    name="unused_defs_tool",
                    description="Tool with $defs but no $ref usage",
                    parameters={
                        "type": "object",
                        "properties": {
                            "data": {"type": "string"},
                        },
                        "required": ["data"],
                        "$defs": {
                            "UnusedType": {
                                "type": "object",
                                "properties": {
                                    "value": {"type": "string"},
                                },
                            },
                        },
                    },
                ),
            ),
        ]

        try:
            _get_tool_schema_defs(tools_with_unused_defs)
        except ValueError as e:
            self.fail(f"Should not raise ValueError, but got: {e}")

        schema = get_json_schema_constraint(tools_with_unused_defs, "required")
        self.assertIsNotNone(schema)
        jsonschema.Draft202012Validator.check_schema(schema)

        # Should still include $defs even if not referenced
        self.assertIn("$defs", schema)
        self.assertIn("UnusedType", schema["$defs"])


if __name__ == "__main__":
    unittest.main()