from typing import Iterator
import re
import xml.etree.ElementTree as ET

from .inference_result import InferenceResult

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)

def stop_before_value(iterator: Iterator[str], stop_value: str) -> Iterator[str]:
    """
    Creates an iterator that yields values from the input iterator
    until it encounters the stop_value (exclusive).
    
    Args:
        iterator: The source iterator
        stop_value: The value to stop before
        
    Yields:
        Values from the iterator until stop_value is encountered
        If stop_value is part of an item, yields the part before stop_value
    """
    for item in iterator:
        if stop_value in item:
            split_point = item.index(stop_value)
            if split_point > 0:
                yield item[:split_point]
            break
        yield item