# -*- coding: utf8 -*-

import xml.etree.ElementTree as ET
from types import StringType, UnicodeType
import biocase.configuration
cfg=biocase.configuration.Cfg()

import logging
log = logging.getLogger("lib.tools.xml")


def localName(elem):
    if type(elem.tag) not in [StringType, UnicodeType]:
        return ""
    if elem.tag[0] != "{":
        tag = elem.tag
    uri, tag = elem.tag[1:].split("}")
    return tag
    
def indent(elem, level=0):
    i = "\n" + level*"  "
    if len(elem):
        if not elem.text or not elem.text.strip():
            elem.text = i + "  "
        for e in elem:
            indent(e, level+1)
        if not e.tail or not e.tail.strip():
            e.tail = i
    if level and (not elem.tail or not elem.tail.strip()):
        elem.tail = i


def prettyprint(elem, stream=None, encoding='utf-8'):
    indent(elem)
    if stream is None:
        return ET.tostring(elem, encoding)
    else:
        stream.write(ET.tostring(elem, encoding))
    
def element(ns, tagname, content=None, attrNs=None, **attrs):
    node = ET.Element("{%s}%s"%(ns,tagname))
    addOptionalAttributes(node, attrNs, **attrs)
    if content is not None:
        node.text=unicode(content)
    return node

def subelement(parentelem, ns, tagname, content=None, attrNs=None, **attrs):
    node = element(ns=ns, tagname=tagname, content=content, attrNs=attrNs, **attrs)
    parentelem.append(node)
    return node

def addOptionalAttributes(elem, ns, **attrs):
    for k,v in attrs.items():
        if v is not None and len(str(v))>0:
            if ns is None:
                elem.set('%s'%(k), unicode(v))
            else:
                elem.set('{%s}%s'%(ns, k), unicode(v))
    return elem

def clone(elem, deep=True):
    attrs=dict([(a,elem.get(a)) for a in elem.keys()])
    copy=elem.makeelement(elem.tag, attrs)
    copy.text=elem.text
    if deep:
        for e in elem:
            copy.append(clone(e, True))
    return copy

def cloneAlt(elem, deep=True):
    copy=ET.Element(elem.tag)
    copy.text=elem.text
    for a in elem.keys():
        copy.set(a, elem.get(a))
    if deep:
        for e in elem:
            copy.append(cloneAlt(e, True))
    return copy
    
def set_prefixes(elem, prefix_map):
    '''Here’s a simple function that takes a prefix/uri mapping, creates the necessary xmlns attributes on the given tree, and changes all universal names to their equivalent prefix:tag form'''
    # check if this is a tree wrapper
    if not ET.iselement(elem):
        elem = elem.getroot()
    # build uri map and add to root element
    uri_map = {}
    for prefix, uri in prefix_map.items():
        uri_map[uri] = prefix
        elem.set("xmlns:" + prefix, uri)
    # fixup all elements in the tree
    memo = {}
    for elem in elem.getiterator():
        # Comments are ElementTree instances but have no tag string attribute!
        if type(elem.tag) in [StringType, UnicodeType]:
            fixup_element_prefixes(elem, uri_map, memo)

def fixup_element_prefixes(elem, uri_map, memo):
    '''This checks all universal names against the URI map, for both element and attribute names.'''
    def fixup(name):
        try:
            return memo[name]
        except KeyError:
            if name[0] != "{":
                return
            uri, tag = name[1:].split("}")
            if uri in uri_map:
                new_name = uri_map[uri] + ":" + tag
                memo[name] = new_name
                return new_name
    # fix element name
    name = fixup(elem.tag)
    if name:
        elem.tag = name
    # fix attribute names
    for key, value in elem.items():
        name = fixup(key)
        if name:
            elem.set(name, value)
            del elem.attrib[key]

def cleanupNamespaces(tree):
    # get unique namespace URI:s
    log.debug("Cleanup namespaces for %s"%tree)
    uri_set = set()
    for elem in tree.getiterator():
        # Comments are ElementTree instances but have no tag string attribute!
        if type(elem.tag) in [StringType, UnicodeType]:
            if hasattr(elem,"tag") and elem.tag[0] == "{":
                uri, tag = elem.tag[1:].split("}")
                uri_set.add(uri)
            # check attributes too
            for attr in elem.keys():
                if attr[0] == "{":
                    uri, tag = attr[1:].split("}")
                    uri_set.add(uri)
    # create prefix map
    prefix_map = {}
    log.debug("Unique namespaces found in response message to be prefixed: %s"%str(uri_set))
    for uri in sorted(uri_set):
        if cfg.ns2prefix.has_key(uri):
            prefix_map[cfg.ns2prefix[uri]] = uri
        else:
            prefix_map["ns%d" % len(prefix_map)] = uri
    log.info("Namespace prefixes to be assigned: %s"%str(prefix_map))
    set_prefixes(tree, prefix_map)


def countElementsOfPath(node, path):
    '''takes a ElementTree node and a simple xpath expression and counts all nodes
    in the object that belongs to the xpath expression.
    '''
    if len(path) == 0 or node == None:
        return 0
    #    
    if path[0] == '/':
        # this path still has the root slash. remove it.
        path = path[1:]
    if path.find('/') > 0:
        # slash found. we have to check for children
        nextPathPart = path.split('/')[0]
        path = path[len(nextPathPart)+1:]
        hits = 0
        if localName(node) == nextPathPart:
            for c in node:
                hits += countElementsOfPath(c, path)
        return hits
    else:
        # cut off possible attributes
        if path.find('@') > 0:
            path = path[:path.index('@')]
        # is this the wanted node ?
        if localName(node) == path:
            return 1
        else:
            return 0




if __name__ ==  "__main__":
    ns = "http://www.biocase.org/schemas/protocol/1.3"
    loc = "http://www.biocase.org/protocol/1.3/protocol.xsd"
    ns2="http://www.biocase.org/rules"
    xsi = 'http://www.w3.org/2001/XMLSchema-instance'

    psf_node = element(ns, "Provider", adminpassword="adminpassword")
    attrDict={"heart":"pia","love":"pia"}
    addOptionalAttributes(psf_node, xsi, schemaLocation=ns+' '+loc, **attrDict)
    c = ET.Comment("Tachchen Alter")
    #psf_node[0:0]=[c]
    psf_node.append(c)
    host = subelement(psf_node, ns, "Host")
    subelement(host, ns, "User", "user")
    subelement(host, ns, "Password", "password")
    subelement(host, ns, "Datasource", "datasource")
    subelement(host, ns, "DB_IP", "dbIP")
    subelement(host, ns, "DBMS", "dbms")
    db_node = subelement(psf_node, ns, "DB")
    subelement(db_node, ns, "TextEncoding", "encoding")
    host[1:1]=[element(ns2, "destination","192.534.25.4"),element(ns2, "destination","192.534.25.5")]
    print prettyprint(psf_node)
    print "-"*40
    print countElementsOfPath(psf_node, "/Provider/Host/destination")
    print psf_node.keys()
    print "-"*40
    cleanupNamespaces(psf_node)
    print prettyprint(psf_node)
    print "-"*40
    copy=clone(psf_node)
    print prettyprint(copy)
    
    print "-"*40
    psf = element(ns, "Provider")
    host = subelement(psf, ns, "Host", "my BODY", adminpassword="cfhuwicheuwba", karla="klug")
    host = subelement(psf, ns, "Horst")
    host = subelement(psf, ns, "Harry", look=None, large="xxxl", club=42)
    host = subelement(psf, ns2, "Home", None)
    host = subelement(psf, ns2, "Homie", "")
    host = subelement(psf, ns, "Harn", 23)
    print prettyprint(psf)
    print "-"*40
    cleanupNamespaces(psf)
    print prettyprint(psf)
    
    print c
    print host
    print callable(host)
    print type(c.tag)
    print type(host.tag) in [StringType, UnicodeType]
    print type(c.tag) in [StringType, UnicodeType]
    