from rdflib.serializer import Serializer
from rdflib.plugins.serializers.xmlwriter import XMLWriter

from rdflib.term import URIRef, Literal, BNode
from rdflib.namespace import Namespace

from rdflib.graph import Graph, ConjunctiveGraph

from rdflib.py3compat import b

__all__ = ['TriXSerializer']

## TODO: MOve this somewhere central
TRIXNS=Namespace("http://www.w3.org/2004/03/trix/trix-1/")
XMLNS=Namespace("http://www.w3.org/XML/1998/namespace")

class TriXSerializer(Serializer):
    def __init__(self, store):
        super(TriXSerializer, self).__init__(store)
        if not store.context_aware: 
            raise Exception("TriX serialization only makes sense for context-aware stores!")


    def serialize(self, stream, base=None, encoding=None, **args):

        nm=self.store.namespace_manager 

        self.writer=XMLWriter(stream, nm, encoding, extra_ns={"": TRIXNS})

        self.writer.push(TRIXNS["TriX"])
        self.writer.namespaces()

        if isinstance(self.store, ConjunctiveGraph):
            for subgraph in self.store.contexts():
                self._writeGraph(subgraph)
        elif isinstance(self.store, Graph):
            self._writeGraph(self.store)
        else:
            raise Exception("Unknown graph type: "+type(self.store))

        self.writer.pop()
        stream.write(b("\n"))
        

    def _writeGraph(self, graph):
        self.writer.push(TRIXNS["graph"])
        if isinstance(graph.identifier, URIRef):
            self.writer.element(TRIXNS["uri"], content=str(graph.identifier))
            
        for triple in graph.triples((None,None,None)):
            self._writeTriple(triple)
        self.writer.pop()

    def _writeTriple(self, triple):
        self.writer.push(TRIXNS["triple"])
        for component in triple:
            if isinstance(component, URIRef):
                self.writer.element(TRIXNS["uri"],
                                                                    content=str(component))
            elif isinstance(component, BNode):
                self.writer.element(TRIXNS["id"],
                                                                    content=str(component))
            elif isinstance(component, Literal):
                if component.datatype:
                    self.writer.element(TRIXNS["typedLiteral"],
                        content=str(component),
                        attributes={ TRIXNS["datatype"]: str(component.datatype) })
                elif component.language:
                    self.writer.element(TRIXNS["plainLiteral"],
                        content=str(component),
                        attributes={ XMLNS["lang"]: str(component.language) })
                else:
                    self.writer.element(TRIXNS["plainLiteral"],
                        content=str(component))
        self.writer.pop()
