#include "parse_event_handler.h"
#include "expat_module.h"
#include "domlette.h"

#define INITIAL_CHILDREN 4

typedef struct _context {
  struct _context *next;
  PyNodeObject *node;

  PyObject *xml_base;

  PyNodeObject **children;
  int children_allocated;
  int children_size;
} Context;

typedef struct {
  ExpatParser parser;
  PyDocumentObject *owner_document;

  Context *context;
  Context *free_context;

  PyObject *new_namespaces;
} ParserState;

static PyObject *xmlns_string;
static PyObject *base_string;
static PyObject *process_includes_string;
static PyObject *strip_elements_string;
static PyObject *empty_args_tuple;
static PyObject *gc_enable_function;
static PyObject *gc_disable_function;
static PyObject *gc_isenabled_function;

/** Utility Functions **************************************************/

static PyObject *MakeQName(PyObject *prefix, PyObject *localName)
{
  PyObject *nodeName;

  if (PyObject_IsTrue(prefix)) {
    Py_UNICODE *ustr;

    nodeName = PyUnicode_FromUnicode(NULL,
                                     PyUnicode_GET_SIZE(prefix) + \
                                     PyUnicode_GET_SIZE(localName) + 1);
    if (nodeName == NULL) return NULL;

    ustr = PyUnicode_AS_UNICODE(nodeName);
    Py_UNICODE_COPY(ustr,
                    PyUnicode_AS_UNICODE(prefix),
                    PyUnicode_GET_SIZE(prefix));

    ustr[PyUnicode_GET_SIZE(prefix)] = (Py_UNICODE)':';

    Py_UNICODE_COPY(ustr + PyUnicode_GET_SIZE(prefix) + 1,
                    PyUnicode_AS_UNICODE(localName),
                    PyUnicode_GET_SIZE(localName));
  } else {
    Py_INCREF(localName);
    nodeName = localName;
  }
  return nodeName;
}


/** Context ************************************************************/

static Context *Context_New(void)
{
  Context *self = PyMem_New(Context, 1);
  if (self == NULL) {
    PyErr_NoMemory();
    return NULL;
  }
  memset(self, 0, sizeof(Context));

  self->children = PyMem_New(PyNodeObject *, INITIAL_CHILDREN);
  if (self->children == NULL) {
    PyErr_NoMemory();
    PyMem_Free(self);
    return NULL;
  }
  self->children_allocated = INITIAL_CHILDREN;

  return self;
}


static void Context_Del(Context *self)
{
  int i;

  if (self->next) {
    Context_Del(self->next);
  }

  for (i = self->children_size; --i > 0;) {
    Py_DECREF(self->children[i]);
  }
  PyMem_Free(self->children);

  PyMem_Free(self);
}


/** ParserState ********************************************************/


static ParserState *ParserState_New(PyDocumentObject *doc,
                                    PyNodeObject *rootNode)
{
  ParserState *self = PyMem_New(ParserState, 1);
  if (self == NULL) {
    PyErr_NoMemory();
    return NULL;
  }

  /* context */
  if ((self->context = Context_New()) == NULL) {
    PyMem_Free(self);
    return NULL;
  }
  self->context->node = rootNode;
  self->context->xml_base = PyDocument_BASE_URI(doc);
  self->free_context = NULL;
  
  self->new_namespaces = NULL;

  Py_INCREF(doc);
  self->owner_document = doc;

  return self;
}


static void ParserState_Del(ParserState *self)
{
  if (self->context)
    Context_Del(self->context);

  if (self->free_context)
    Context_Del(self->free_context);

  if (self->new_namespaces) {
    Py_DECREF(self->new_namespaces);
  }

  Py_DECREF(self->owner_document);

  PyMem_Free(self);
}


static Context *ParserState_GetContext(ParserState *self)
{
  Context *context = self->free_context;
  if (context != NULL) {
    /* reuse an existing context */
    self->free_context = context->next;
  } else {
    /* create a new context */
    context = Context_New();
    if (context == NULL) return NULL;
  }

  /* make it the active context */
  context->next = self->context;
  self->context = context;
  return context;
}


static void ParserState_FreeContext(ParserState *self)
{
  Context *context = self->context;
  if (context != NULL) {
    /* switch the active context to the following one */
    self->context = context->next;

    /* move this one to the free list */
    context->next = self->free_context;
    self->free_context = context;

    /* reset its values */
    context->node = NULL;
    context->xml_base = NULL;
    context->children_size = 0;
  }
}


static int ParserState_AddNode(ParserState *self, PyNodeObject *node)
{
  Context *context = self->context;
  PyNodeObject **children = context->children;
  int newsize = context->children_size + 1;

  if (node == NULL) {
    PyErr_BadInternalCall();
    return 0;
  }

  /* increase size of child array, if needed */
  if (newsize >= context->children_allocated) {
    int new_allocated = newsize << 1;
    if (PyMem_Resize(children, PyNodeObject *, new_allocated) == NULL) {
      PyErr_NoMemory();
      return 0;
    }
    context->children = children;
    context->children_allocated = new_allocated;
  }

  /* add the node to the children array */
  children[context->children_size] = node;
  context->children_size = newsize;
  return 1;
}


/** handlers ***********************************************************/


static void builder_EndDocument(void *userState)
{
  ParserState *state = (ParserState *)userState;
  Context *context = state->context;

  /* Set the document's children */
  if (_Node_SetChildren(context->node, context->children,
                        context->children_size) == -1) {
    Expat_ParserStop(state->parser);
    return;
  }

  /* Mark the current context as free */
  ParserState_FreeContext(state);
}


static void builder_StartNamespaceDecl(void *userState, PyObject *prefix,
                                       PyObject *uri)
{
  ParserState *state = (ParserState *)userState;

#ifdef DEBUG_PARSER
  fprintf(stderr, "NSDecl started\n");
  fprintf(stderr, "URI:\n");
  _PyObject_Dump(uri);
  fprintf(stderr, "Prefix:\n");
  _PyObject_Dump(prefix);
#endif

  if (state->new_namespaces == NULL) {
    /* first namespace for this element */
    state->new_namespaces = PyDict_New();
    if (state->new_namespaces == NULL) {
      Expat_ParserStop(state->parser);
      return;
    }
  }

  if (uri == Py_None) {
    /* Use an empty string as this will be added as an attribute value.
       Fixes SF#834917
    */
    uri = PyUnicode_FromUnicode(NULL, 0);
    if (uri == NULL) {
      Expat_ParserStop(state->parser);
      return;
    }
  } else {
    Py_INCREF(uri);
  }

  if (PyDict_SetItem(state->new_namespaces, prefix, uri)) {
    Expat_ParserStop(state->parser);
    /* fall through */
  }

  Py_DECREF(uri);
  return;
}

static void builder_StartElement(void *userState, PyObject *name,
                                 PyObject **atts)
{
  ParserState *state = (ParserState *)userState;
  Context *context = state->context;
  PyElementObject *elem = NULL;
  int i;
  PyObject *key, *value;
  PyObject *xml_base;
  PyObject *namespaceURI, *localName, *qualifiedName, *prefix;
  PyObject **ppattr;

#ifdef DEBUG_PARSER
  fprintf(stderr, "Start element\n");
  fprintf(stderr, "name:\n");
  _PyObject_Dump(name);
#endif

  if (Expat_SplitName(state->parser, name, &namespaceURI, &localName, 
                      &qualifiedName, &prefix) == 0) {
    Expat_ParserStop(state->parser);
    return;
  }

  elem = Document_CreateElementNS(state->owner_document, namespaceURI,
                                  qualifiedName, prefix, localName);
  Py_DECREF(namespaceURI);
  Py_DECREF(localName);
  Py_DECREF(qualifiedName);
  Py_DECREF(prefix);
  if (elem == NULL) {
    Expat_ParserStop(state->parser);
    return;
  }

  /** namespaces *******************************************************/

  /* new_namespaces is a dictionary where key is the prefix and value
   * is the uri.
   */
  if (state->new_namespaces) {
    i = 0;
    while (PyDict_Next(state->new_namespaces, &i, &key, &value)) {
      PyAttrObject *nsattr;
      if (key != Py_None) {
        prefix = xmlns_string;
        localName = key;
      } else {
        prefix = key;
        localName = xmlns_string;
      }

      qualifiedName = MakeQName(prefix, localName);
      if (qualifiedName == NULL) {
        Py_DECREF(elem);
        Expat_ParserStop(state->parser);
        return;
      }

      nsattr = Element_SetAttributeNS(elem, g_xmlnsNamespace, qualifiedName, 
                                      prefix, localName, value);
      Py_DECREF(qualifiedName);
      if (nsattr == NULL) {
        Py_DECREF(elem);
        Expat_ParserStop(state->parser);
        return;
      }
      Py_DECREF(nsattr);
    }
    /* make sure children don't set these namespaces */
    Py_DECREF(state->new_namespaces);
    state->new_namespaces = NULL;
  }

  /** attributes *******************************************************/

  xml_base = context->xml_base;

  for (ppattr = atts; *ppattr; ppattr += 2) {
    PyAttrObject *attr;

    if (Expat_SplitName(state->parser, ppattr[0], &namespaceURI, &localName,
                        &qualifiedName, &prefix) == 0) {
      Py_DECREF(elem);
      Expat_ParserStop(state->parser);
      return;
    }

    attr = Element_SetAttributeNS(elem, namespaceURI, qualifiedName, prefix,
                                  localName, ppattr[1]);
    Py_DECREF(namespaceURI);
    Py_DECREF(localName);
    Py_DECREF(qualifiedName);
    Py_DECREF(prefix);
    if (attr == NULL) {
      Py_DECREF(elem);
      Expat_ParserStop(state->parser);
      return;
    }
    Py_DECREF(attr);

    /* check for xml:base */
    if (PyObject_RichCompareBool(namespaceURI, g_xmlNamespace, Py_EQ)) {
      if (PyObject_RichCompareBool(localName, base_string, Py_EQ)) {
        /* use a borrowed value to simplify cleanup */
        xml_base = ppattr[1];
      }
    }
  }

  /* Set XML base */
  Py_DECREF(PyElement_BASE_URI(elem));
  Py_INCREF(xml_base);
  PyElement_BASE_URI(elem) = xml_base;

  /* save states on the context */
  context = ParserState_GetContext(state);
  context->node = (PyNodeObject *) elem;
  context->xml_base = xml_base;
  return;
}


static void builder_EndElement(void *userState, PyObject *name)
{
  ParserState *state = (ParserState *)userState;
  Context *context = state->context;
  PyNodeObject *new_element;

#ifdef DEBUG_PARSER
  fprintf(stderr, "End element\n");
  fprintf(stderr, "name:\n");
  _PyObject_Dump(name);
#endif

  /* Get the newly constructed element */
  new_element = context->node;

  /* Set the element's children */
  if (_Node_SetChildren(new_element, context->children,
                        context->children_size) == -1) {
    Expat_ParserStop(state->parser);
    return;
  }

  /* Mark the current context as free */
  ParserState_FreeContext(state);

  /* ParserState_AddNode steals the reference to the new node */
  if (ParserState_AddNode(state, new_element) == 0) {
    Expat_ParserStop(state->parser);
    /* just fall through */
  }
}


static void builder_CharacterData(void *userState, PyObject *data)
{
  ParserState *state = (ParserState *)userState;
  PyNodeObject *new_text;

#ifdef DEBUG_PARSER
  fprintf(stderr, "Character data:\n");
  _PyObject_Dump(data);
#endif

  new_text = (PyNodeObject *)
    Document_CreateTextNode(state->owner_document, data);

  /* ParserState_AddNode steals the reference to the new node */
  if (ParserState_AddNode(state, new_text) == 0) {
    Expat_ParserStop(state->parser);
    /* just fall through */
  }
}


static void builder_ProcessingInstruction(void *userState, PyObject *target, 
                                          PyObject *data)
{
  ParserState *state = (ParserState *)userState;
  PyNodeObject *new_pi;

  new_pi =(PyNodeObject *)
    Document_CreateProcessingInstruction(state->owner_document, target, data);

  /* ParserState_AddNode steals the reference to the new node */
  if (ParserState_AddNode(state, new_pi) == 0) {
    Expat_ParserStop(state->parser);
    /* just fall through */
  }
}


static void builder_Comment(void *userState, PyObject *data)
{
  ParserState *state = (ParserState *)userState;
  PyNodeObject *new_comment;

  new_comment = (PyNodeObject *)
    Document_CreateComment(state->owner_document, data);

  /* ParserState_AddNode steals the reference to the new node */
  if (ParserState_AddNode(state, new_comment) == 0) {
    Expat_ParserStop(state->parser);
    /* just fall through */
  }
}


static void builder_DoctypeDecl(void *userState, PyObject *systemId, 
                        PyObject *publicId)
{
  ParserState *state = (ParserState *)userState;

  Py_DECREF(state->owner_document->systemId);
  Py_INCREF(systemId);
  state->owner_document->systemId = systemId;

  Py_DECREF(state->owner_document->publicId);
  Py_INCREF(publicId);
  state->owner_document->publicId = publicId;
}


static void builder_UnparsedEntityDecl(void *userState, PyObject *name,
                                       PyObject *uri)
{
  ParserState *state = (ParserState *)userState;

  PyDict_SetItem(state->owner_document->unparsedEntities, name, uri);
}


/** External Routines *************************************************/


static ExpatParser createParser(ParserState *state) {
  ExpatParser parser;

  parser = Expat_ParserCreate(state);
  if (parser == NULL) {
    return NULL;
  }

  Expat_SetEndDocumentHandler(parser, builder_EndDocument);
  Expat_SetStartElementHandler(parser, builder_StartElement);
  Expat_SetEndElementHandler(parser, builder_EndElement);
  Expat_SetStartNamespaceDeclHandler(parser, builder_StartNamespaceDecl);
  Expat_SetCharacterDataHandler(parser, builder_CharacterData);
  Expat_SetProcessingInstructionHandler(parser, builder_ProcessingInstruction);
  Expat_SetCommentHandler(parser, builder_Comment);
  Expat_SetDoctypeDeclHandler(parser, builder_DoctypeDecl);
  Expat_SetUnparsedEntityDeclHandler(parser, builder_UnparsedEntityDecl);

  return parser;
}

PyObject *ParseDocument(PyObject *inputSource, int readDtd, int asEntity)
{
  ParserState *state;
  PyDocumentObject *document;
  PyNodeObject *rootNode;
  PyObject *uri, *strip_elements;
  PyObject *temp;
  int process_xincludes, gc_enabled, rc;

  uri = PyObject_GetAttrString(inputSource, "uri");
  if (uri == NULL) {
    return NULL;
  } else if (!PyUnicode_Check(uri)) {
    PyObject *temp = PyObject_Unicode(uri);
    Py_DECREF(uri);
    if (temp == NULL) return NULL;
    uri = temp;
  }

  document = Document_New(uri);
  Py_DECREF(uri);
  if (document == NULL) {
    return NULL;
  }

  if (asEntity) {
    rootNode = (PyNodeObject *) Document_CreateDocumentFragment(document);
    if (rootNode == NULL) {
      Py_DECREF(document);
      return NULL;
    }
  } else {
    rootNode = (PyNodeObject *) document;
  }

  state = ParserState_New(document, rootNode);
  if (state == NULL) {
    Py_DECREF(document);
    return NULL;
  }

  state->parser = createParser(state);
  if (state->parser == NULL) {
    ParserState_Del(state);
    return NULL;
  }

  /* XSLT whitespace stripping rules */
  strip_elements = PyObject_GetAttr(inputSource, strip_elements_string);
  if (strip_elements == NULL) {
    ParserState_Del(state);
    return NULL;
  }
  if (Expat_SetWhitespaceRules(state->parser, strip_elements) == 0) {
    Py_DECREF(strip_elements);
    ParserState_Del(state);
    return NULL;
  }
  Py_DECREF(strip_elements);
  
  /* Determine whether XInclude processing is enabled */
  temp = PyObject_GetAttr(inputSource, process_includes_string);
  if (temp == NULL) {
    ParserState_Del(state);
    return NULL;
  }
  process_xincludes = PyObject_IsTrue(temp);
  Py_DECREF(temp);
  Expat_SetXIncludeProcessing(state->parser, process_xincludes);

  /* Disable GC (if enabled) while building the DOM tree */
  temp = PyObject_Call(gc_isenabled_function, empty_args_tuple, NULL);
  if (temp == NULL) {
    Expat_ParserFree(state->parser);
    ParserState_Del(state);
    return NULL;
  }
  gc_enabled = PyObject_IsTrue(temp);
  Py_DECREF(temp);
  if (gc_enabled) {
    temp = PyObject_Call(gc_disable_function, empty_args_tuple, NULL);
    if (temp == NULL) {
      Expat_ParserFree(state->parser);
      ParserState_Del(state);
      return NULL;
    }
    Py_DECREF(temp);
  }

  if (asEntity)
    rc = Expat_ParseEntity(state->parser, inputSource);
  else
    rc = Expat_ParseDocument(state->parser, inputSource, readDtd);

  if (gc_enabled) {
    temp = PyObject_Call(gc_enable_function, empty_args_tuple, NULL);
    if (temp == NULL) {
      Expat_ParserFree(state->parser);
      ParserState_Del(state);
      return NULL;
    }
    Py_DECREF(temp);
  }

  Expat_ParserFree(state->parser);
  ParserState_Del(state);
  
  if (rc == 0)
    return NULL;
  else
    return (PyObject *) rootNode;
}

/** Module Setup & Teardown *******************************************/


int DomletteBuilder_Init(PyObject *module)
{
  PyObject *import;

  xmlns_string = PyUnicode_DecodeASCII("xmlns", 5, NULL);
  if (xmlns_string == NULL) return -1;

  base_string = PyUnicode_DecodeASCII("base", 4, NULL);
  if (base_string == NULL) return -1;

  process_includes_string = PyString_FromString("processIncludes");
  if (process_includes_string == NULL) return -1;

  strip_elements_string = PyString_FromString("stripElements");
  if (strip_elements_string == NULL) return -1;

  process_includes_string = PyString_FromString("processIncludes");
  if (process_includes_string == NULL) return -1;

  strip_elements_string = PyString_FromString("stripElements");
  if (strip_elements_string == NULL) return -1;

  empty_args_tuple = PyTuple_New(0);
  if (empty_args_tuple == NULL) return -1;

  import = PyImport_ImportModule("gc");
  if (import == NULL) return -1;

#define GET_GC_FUNC(NAME)                                       \
  gc_##NAME##_function = PyObject_GetAttrString(import, #NAME); \
  if (gc_##NAME##_function == NULL) {                           \
    Py_DECREF(import);                                          \
    return -1;                                                  \
  }

  GET_GC_FUNC(enable);
  GET_GC_FUNC(disable);
  GET_GC_FUNC(isenabled);

  Py_DECREF(import);

  return 0;
}


void DomletteBuilder_Fini(void)
{
  Py_DECREF(xmlns_string);
  Py_DECREF(base_string);
  Py_DECREF(empty_args_tuple);
  Py_DECREF(gc_enable_function);
  Py_DECREF(gc_disable_function);
  Py_DECREF(gc_isenabled_function);
}
