import xml.etree.ElementTree as ET
from typing import Optional, Set

class XMLValidator:
    """
    Validates XML content against a schema.
    
    Attributes:
        _schema: The parsed XML schema to validate against
        _valid_root_elements: Set of valid root element names from schema
    """
    
    def __init__(self, schema: str):
        """
        Initialize validator with XML schema.
        
        Args:
            schema: XML schema string
        """
        # Register namespace used in schema
        ET.register_namespace('xs', 'http://www.w3.org/2001/XMLSchema')
        
        try:
            # Parse schema
            self._schema = ET.fromstring(schema.strip())
            
            # Extract valid root elements
            ns = {'xs': 'http://www.w3.org/2001/XMLSchema'}
            elements = self._schema.findall(".//xs:element", ns)
            self._valid_root_elements = {elem.get('name') for elem in elements if elem.get('name')}
            
        except ET.ParseError as e:
            raise ValueError(f"Invalid schema: {e}")
            
    def validate(self, xml: str) -> Optional[str]:
        """
        Validate XML content against the schema.
        
        Args:
            xml: XML string to validate
            
        Returns:
            str: Error message if validation fails, None if validation succeeds
        """
        try:
            # Parse XML
            root = ET.fromstring(xml.strip())
            
            # Check root element is valid
            if root.tag not in self._valid_root_elements:
                return f"Invalid root element: {root.tag}. Expected one of: {sorted(self._valid_root_elements)}"
                
            # Get schema definition for this element
            ns = {'xs': 'http://www.w3.org/2001/XMLSchema'}
            element_schema = self._schema.find(f".//xs:element[@name='{root.tag}']", ns)
            if element_schema is None:
                return f"Schema definition not found for element: {root.tag}"
                
            # Validate attributes if complex type defined
            complex_type = element_schema.find('xs:complexType', ns)
            if complex_type is not None:
                # Check required attributes
                for attr in complex_type.findall('.//xs:attribute[@use="required"]', ns):
                    attr_name = attr.get('name')
                    if attr_name not in root.attrib:
                        return f"Missing required attribute '{attr_name}' on element '{root.tag}'"
                        
                # Check attribute types
                for attr_name, attr_value in root.attrib.items():
                    attr_schema = complex_type.find(f'.//xs:attribute[@name="{attr_name}"]', ns)
                    if attr_schema is None:
                        return f"Unexpected attribute '{attr_name}' on element '{root.tag}'"
                        
                    attr_type = attr_schema.get('type')
                    if attr_type == 'xs:string':
                        continue  # All string values are valid
                    elif attr_type == 'xs:integer':
                        try:
                            int(attr_value)
                        except ValueError:
                            return f"Invalid integer value '{attr_value}' for attribute '{attr_name}'"
                            
            return None  # Validation successful
            
        except ET.ParseError as e:
            return f"Invalid XML: {e}"
        except Exception as e:
            return f"Validation error: {e}"
            
    def get_valid_root_elements(self) -> Set[str]:
        """
        Get set of valid root element names from schema.
        
        Returns:
            Set[str]: Set of valid root element names
        """
        return self._valid_root_elements.copy()