from string import Template
# Minimal stubs for required classes and enums to support the function.
from typing import Any, Dict, List, Optional, Union

# imports
import pytest
from guardrails.schema.rail_schema import rail_string_to_schema
# Begin function definition
from lxml import etree as ET
from lxml.etree import XMLParser, _Element


class SimpleTypes:
    STRING = "string"
    INTEGER = "integer"
    NUMBER = "number"
    BOOLEAN = "boolean"
    OBJECT = "object"
    ARRAY = "array"

class OutputTypes:
    STRING = "string"
    LIST = "list"
    DICT = "dict"

class GuardExecutionOptions:
    def __init__(self):
        self.messages = []
        self.reask_messages = []

class ProcessedSchema:
    def __init__(
        self,
        validators: Optional[List[Any]] = None,
        validator_map: Optional[Dict[str, Any]] = None,
        exec_opts: Optional[GuardExecutionOptions] = None,
    ):
        self.validators = validators if validators is not None else []
        self.validator_map = validator_map if validator_map is not None else {}
        self.exec_opts = exec_opts if exec_opts is not None else GuardExecutionOptions()
        self.json_schema = None
        self.output_type = None
from guardrails.schema.rail_schema import rail_string_to_schema

# ---------------------- UNIT TESTS ----------------------

# Basic Test Cases

def test_basic_string_output():
    """Test simple string output schema."""
    rail = """<rail>
        <output type="string" description="A simple string." />
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 146μs -> 116μs (26.3% faster)

def test_basic_object_output():
    """Test simple object output schema with one field."""
    rail = """<rail>
        <output type="object" description="An object.">
            <string name="foo" description="A foo string." />
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 190μs -> 145μs (31.1% faster)

def test_basic_list_output():
    """Test simple list output schema."""
    rail = """<rail>
        <output type="list" description="A list of integers.">
            <integer description="An integer item." />
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 185μs -> 139μs (33.0% faster)

def test_basic_enum_output():
    """Test enum output schema."""
    rail = """<rail>
        <output type="object">
            <enum name="color" values="red, green, blue" description="Color." />
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 193μs -> 147μs (30.9% faster)
    props = schema.json_schema["properties"]

def test_basic_messages_and_reask_messages():
    """Test that messages and reask_messages are extracted."""
    rail = """<rail>
        <output type="string" />
        <messages>
            <message role="system">Hello!</message>
            <message role="user">Hi!</message>
        </messages>
        <reask_messages>
            <message role="system">Try again.</message>
        </reask_messages>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 176μs -> 131μs (34.0% faster)

# Edge Test Cases

def test_missing_output_tag_raises():
    """Test that missing <output> raises ValueError."""
    rail = """<rail>
        <messages>
            <message role="system">Hello!</message>
        </messages>
    </rail>"""
    with pytest.raises(ValueError, match="RAIL must contain a output element!"):
        rail_string_to_schema(rail) # 37.0μs -> 35.2μs (5.10% faster)


def test_list_with_multiple_children_raises():
    """Test that <list> with multiple children raises ValueError."""
    rail = """<rail>
        <output type="list">
            <string />
            <integer />
        </output>
    </rail>"""
    with pytest.raises(ValueError, match="<list /> RAIL elements must have precisely 1 child element!"):
        rail_string_to_schema(rail) # 77.7μs -> 75.7μs (2.73% faster)



def test_object_with_no_children_additional_properties_true():
    """Test that object with no children sets additionalProperties to True."""
    rail = """<rail>
        <output type="object" description="Empty object." />
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 158μs -> 128μs (23.1% faster)

def test_enum_with_empty_values():
    """Test enum with empty values attribute."""
    rail = """<rail>
        <output type="object">
            <enum name="empty_enum" values="" description="Empty enum." />
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 192μs -> 145μs (32.5% faster)

def test_list_with_no_child_items_empty_dict():
    """Test that <list> with no child gives items={}."""
    rail = """<rail>
        <output type="list" />
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 142μs -> 110μs (28.4% faster)

def test_object_with_optional_field():
    """Test object with required and optional fields."""
    rail = """<rail>
        <output type="object">
            <string name="foo" required="true" />
            <integer name="bar" required="false" />
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 226μs -> 165μs (36.4% faster)

def test_string_with_format():
    """Test <string> with format attribute."""
    rail = """<rail>
        <output type="object">
            <string name="email" format="email" />
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 186μs -> 143μs (30.2% faster)

def test_date_with_custom_format():
    """Test <date> with custom format and date-format attributes."""
    rail = """<rail>
        <output type="object">
            <date name="dob" format="custom" date-format="%Y-%m-%d" />
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 196μs -> 152μs (29.5% faster)

# Large Scale Test Cases

def test_large_object_schema():
    """Test large object with many fields."""
    num_fields = 500
    fields = "\n".join(
        [f'<string name="field{i}" description="Field {i}" />' for i in range(num_fields)]
    )
    rail = f"""<rail>
        <output type="object">
            {fields}
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 14.5ms -> 7.85ms (84.8% faster)
    props = schema.json_schema["properties"]
    for i in range(num_fields):
        pass

def test_large_list_schema():
    """Test large list with a complex item schema."""
    rail = """<rail>
        <output type="list">
            <object>
                {}
            </object>
        </output>
    </rail>""".format(
        "\n".join([f'<integer name="num{i}" />' for i in range(100)])
    )
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 3.03ms -> 1.71ms (76.9% faster)
    items = schema.json_schema["items"]
    for i in range(100):
        pass

def test_large_enum_schema():
    """Test enum with many values."""
    values = ",".join([f"v{i}" for i in range(200)])
    rail = f"""<rail>
        <output type="object">
            <enum name="big_enum" values="{values}" />
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 223μs -> 179μs (24.6% faster)
    enum_values = schema.json_schema["properties"]["big_enum"]["enum"]
    for i in range(200):
        pass


def test_large_nested_object():
    """Test deeply nested object."""
    rail = """<rail>
        <output type="object">
            <object name="level1">
                <object name="level2">
                    <object name="level3">
                        <string name="deep_field" />
                    </object>
                </object>
            </object>
        </output>
    </rail>"""
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 284μs -> 210μs (35.1% faster)
    props = schema.json_schema["properties"]
    l1 = props["level1"]["properties"]
    l2 = l1["level2"]["properties"]
    l3 = l2["level3"]["properties"]
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import xml.etree.ElementTree as ET

# imports
import pytest
from guardrails.schema.rail_schema import rail_string_to_schema


# Minimal stubs for dependencies used in rail_string_to_schema
class SimpleTypes:
    STRING = "string"
    INTEGER = "integer"
    NUMBER = "number"
    BOOLEAN = "boolean"
    OBJECT = "object"
    ARRAY = "array"

class OutputTypes:
    STRING = "string"
    LIST = "list"
    DICT = "dict"

class GuardExecutionOptions:
    def __init__(self):
        self.messages = None
        self.reask_messages = None

class ProcessedSchema:
    def __init__(self, validators, validator_map, exec_opts):
        self.validators = validators
        self.validator_map = validator_map
        self.exec_opts = exec_opts
        self.json_schema = None
        self.output_type = None
from guardrails.schema.rail_schema import rail_string_to_schema

# ---------------------- UNIT TESTS ----------------------

# Basic Test Cases
def test_basic_string_output():
    # Test with a simple string output
    rail = """
    <rail>
      <output type="string" description="A simple string output." />
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 142μs -> 112μs (26.4% faster)

def test_basic_object_output():
    # Test with a simple object output with two fields
    rail = """
    <rail>
      <output type="object" description="User info.">
        <string name="username" description="User's name." />
        <integer name="age" description="User's age." />
      </output>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 228μs -> 168μs (35.7% faster)

def test_basic_list_output():
    # Test with a simple list output of strings
    rail = """
    <rail>
      <output type="list" description="List of tags.">
        <string description="A tag." />
      </output>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 182μs -> 139μs (30.5% faster)

def test_basic_enum_output():
    # Test with an enum output
    rail = """
    <rail>
      <output type="enum" description="Color">
        <enum values="red, green, blue" description="Color enum." />
      </output>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 145μs -> 116μs (24.9% faster)
    # The values should be ['red', 'green', 'blue']

def test_basic_messages_and_reask_messages():
    # Test with messages and reask_messages
    rail = """
    <rail>
      <output type="string" />
      <messages>
        <message role="system">Hello!</message>
        <message role="user">What's your name?</message>
      </messages>
      <reask_messages>
        <message role="system">Try again.</message>
      </reask_messages>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 178μs -> 133μs (33.9% faster)

# Edge Test Cases

def test_missing_output_element():
    # Should raise ValueError if output element is missing
    rail = """
    <rail>
      <messages>
        <message role="system">No output!</message>
      </messages>
    </rail>
    """
    with pytest.raises(ValueError, match="RAIL must contain a output element!"):
        rail_string_to_schema(rail) # 36.9μs -> 36.2μs (1.92% faster)

def test_list_with_multiple_children():
    # Should raise ValueError if <list> has more than one child
    rail = """
    <rail>
      <output type="list">
        <string />
        <integer />
      </output>
    </rail>
    """
    with pytest.raises(ValueError, match="must have precisely 1 child"):
        rail_string_to_schema(rail) # 73.3μs -> 71.2μs (2.94% faster)

def test_choice_missing_discriminator():
    # Should raise ValueError if <choice> is missing discriminator
    rail = """
    <rail>
      <output type="choice">
        <case name="A">
          <string name="foo" />
        </case>
      </output>
    </rail>
    """
    with pytest.raises(ValueError, match="must specify a discriminator"):
        rail_string_to_schema(rail) # 73.5μs -> 71.4μs (2.91% faster)

def test_case_missing_name():
    # Should raise ValueError if <case> is missing name
    rail = """
    <rail>
      <output type="choice" discriminator="type">
        <case>
          <string name="foo" />
        </case>
      </output>
    </rail>
    """
    with pytest.raises(ValueError, match="must specify a name"):
        rail_string_to_schema(rail) # 116μs -> 87.5μs (33.1% faster)

def test_object_with_nameless_child():
    # Should skip nameless child and not crash
    rail = """
    <rail>
      <output type="object">
        <string name="foo" />
        <integer />
      </output>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 702μs -> 653μs (7.58% faster)

def test_object_with_no_children():
    # Should set additionalProperties=True if no children
    rail = """
    <rail>
      <output type="object" />
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 153μs -> 121μs (25.7% faster)

def test_list_with_no_children():
    # Should set items to {} if no children
    rail = """
    <rail>
      <output type="list" />
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 143μs -> 112μs (27.3% faster)

def test_enum_with_no_values():
    # Should set enum to None or empty
    rail = """
    <rail>
      <output type="enum">
        <enum />
      </output>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 141μs -> 112μs (25.5% faster)

def test_custom_format_handling():
    # Should handle custom format attributes correctly
    rail = """
    <rail>
      <output type="date" description="A date." format="foo" date-format="%Y-%m-%d" />
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 153μs -> 122μs (24.7% faster)

def test_required_false_handling():
    # Should not include field in required if required="false"
    rail = """
    <rail>
      <output type="object">
        <string name="foo" required="false" />
        <string name="bar" required="true" />
      </output>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 226μs -> 162μs (38.9% faster)

# Large Scale Test Cases

def test_large_object_output():
    # Test with a large object output
    num_fields = 100
    fields = "\n".join([f'<string name="field{i}" />' for i in range(num_fields)])
    rail = f"""
    <rail>
      <output type="object">
        {fields}
      </output>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 2.94ms -> 1.62ms (82.1% faster)
    for i in range(num_fields):
        pass

def test_large_list_output():
    # Test with a large list output of objects
    rail = """
    <rail>
      <output type="list">
        <object>
          {}
        </object>
      </output>
    </rail>
    """.format("\n".join([f'<integer name="num{i}" />' for i in range(50)]))
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 1.63ms -> 940μs (73.5% faster)
    for i in range(50):
        pass

def test_large_enum_output():
    # Test with a large enum output
    enum_values = ",".join([f"val{i}" for i in range(200)])
    rail = f"""
    <rail>
      <output type="enum">
        <enum values="{enum_values}" />
      </output>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 146μs -> 117μs (24.8% faster)
    for i in range(200):
        pass

def test_large_choice_output():
    # Test with a choice output with many cases
    cases = "\n".join([
        f'<case name="case{i}"><string name="field{i}" /></case>'
        for i in range(20)
    ])
    rail = f"""
    <rail>
      <output type="choice" discriminator="type">
        {cases}
      </output>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 1.48ms -> 1.08ms (36.6% faster)
    for i in range(20):
        pass

def test_large_messages():
    # Test with many messages
    messages = "\n".join([f'<message role="user">msg{i}</message>' for i in range(50)])
    rail = f"""
    <rail>
      <output type="string" />
      <messages>
        {messages}
      </messages>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 217μs -> 168μs (29.1% faster)
    for i in range(50):
        pass

def test_large_reask_messages():
    # Test with many reask messages
    messages = "\n".join([f'<message role="system">reask{i}</message>' for i in range(50)])
    rail = f"""
    <rail>
      <output type="string" />
      <reask_messages>
        {messages}
      </reask_messages>
    </rail>
    """
    codeflash_output = rail_string_to_schema(rail); schema = codeflash_output # 184μs -> 155μs (19.2% faster)
    for i in range(50):
        pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.