from .inference_result import InferenceResult

import xml.etree.ElementTree as ET
import re

def get_valid_root_elements(schema: str) -> set:
    """
    Extract valid root element names from the XML schema.
    
    Args:
        schema: XML schema string
        
    Returns:
        set: Set of valid root element names
    """
    try:
        schema = schema.strip()
        ET.register_namespace('xs', 'http://www.w3.org/2001/XMLSchema')
        root = ET.fromstring(schema)
        ns = {'xs': 'http://www.w3.org/2001/XMLSchema'}
        elements = root.findall(".//xs:element", ns)
        return {elem.get('name') for elem in elements if elem.get('name')}
    except ET.ParseError as e:
        print(f"Error parsing schema: {e}")
        return set()

def split_response(response: str, valid_elements: set) -> InferenceResult:
    """
    Split the response into reasoning and actions based on valid XML elements.
    
    Args:
        response: Raw response string from the model
        valid_elements: Set of valid root element names from the schema
        
    Returns:
        InferenceResult: Tuple containing reasoning and actions
    """
    elements_pattern = '|'.join(map(re.escape, valid_elements))
    pattern = f"<({elements_pattern})[^>]*>"
    matches = list(re.finditer(pattern, response))
    if not matches:
        return InferenceResult(response.strip(), "")
    last_match = matches[-1]
    split_point = last_match.start()
    reasoning = response[:split_point].strip()
    actions = response[split_point:].strip()
    return InferenceResult(reasoning, actions)