from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import xml.etree.ElementTree as ET
from compas import _iotools
from compas.files._xml.xml_cpython import prettify_string  # doesn't need special handling for pre-3.8 so we just import

__all__ = [
    'xml_from_file',
    'xml_from_string',
    'prettify_string',
]


def xml_from_file(source, tree_parser=None):
    if tree_parser:
        raise NotImplementedError('XML parsing on CPython 3.7 and older does not support a custom tree parser')

    tree_parser = ET.XMLPullParser
    parser = tree_parser(events=('start', 'start-ns'))

    with _iotools.open_file(source) as file:
        for data in _iotools.iter_file(file):
            parser.feed(data)
        return _process_all_events(parser)


def xml_from_string(text, tree_parser=None):
    if tree_parser:
        raise NotImplementedError('XML parsing on CPython 3.7 and older does not support a custom tree parser')

    tree_parser = ET.XMLPullParser
    parser = tree_parser(events=('start', 'end', 'start-ns', 'end-ns'))
    parser.feed(text)
    return _process_all_events(parser)


def _process_all_events(parser):
    root = None
    current_namespaces = {}

    for event, event_data in parser.read_events():
        if event == 'start':
            element = event_data
            if not root:
                root = element

            if len(current_namespaces):
                element.attrib.update(current_namespaces)

            current_namespaces = {}

        if event == 'start-ns':
            prefix, uri = event_data
            ns_prefix = 'xmlns:' + prefix if prefix else 'xmlns'
            current_namespaces[ns_prefix] = uri

    parser.close()

    return root
