# coding=utf-8
"""
Jinja2 extensions
"""
from __future__ import absolute_import

from functools import partial

from jinja2.ext import Extension
from jinja2 import nodes

from werkzeug.local import LocalProxy
from flask import current_app
from flask.globals import _request_ctx_stack, _lookup_req_object

deferred_js = LocalProxy(partial(_lookup_req_object, 'deferred_js'))

class DeferredJS(object):
  """
  Flask extentions for use with DeferredJSExtension for jinja
  """
  name = 'deferred_js'

  def __init__(self, app=None):
    if app is not None:
      self.init_app(app)

  def init_app(self, app):
    if self.name in app.extensions:
      return

    app.extensions[self.name] = self
    app.before_request(self.reset_deferred)

  def reset_deferred(self):
    _request_ctx_stack.top.deferred_js = []


class DeferredJSExtension(Extension):
  """
  Put JS fragment at the end of the document in a script tag.
  """
  tags = set(['deferJS', 'deferredJS'])

  def parse(self, parser):
    token = next(parser.stream)
    tag = token.value
    lineno = token.lineno

    # now we parse body of the block
    body = parser.parse_statements(['name:enddeferJS', 'name:enddeferredJS'],
                                   drop_needle=True)

    method = 'defer_nodes' if tag == 'deferJS' else 'collect_deferred'
    return nodes.CallBlock(self.call_method(method, []),
                           [], [], body).set_lineno(lineno)

  def defer_nodes(self, caller):
    body = caller()
    deferred_js.append(body)
    return u''

  def collect_deferred(self, caller):
    result = u'\n'.join(
      u'{{\n{}\n}};'.format(js)
      for js in deferred_js)
    current_app.extensions[DeferredJS.name].reset_deferred()
    return result

@@ -105,16 +105,33 @@
 artial%0A%0A
+import lxml.html%0A
 from jin
@@ -929,16 +929,147 @@
 ipt tag.
+%0A%0A  The JS fragment can contains %3Cscript%3E tag so that your favorite editor%0A  keeps doing proper indentation, syntax highlighting...
 %0A  %22%22%22%0A 
@@ -1622,16 +1622,453 @@
 y = 
-caller()
+u'%3Cdiv%3E%7B%7D%3C/div%3E'.format(caller().strip())%0A%0A    # remove 'script' tag in immediate children, if any%0A    fragment = lxml.html.fragment_fromstring(body)%0A    for child in fragment:%0A      if child.tag == 'script':%0A        child.drop_tag() # side effect on fragment.text or previous_child.tail!%0A%0A    body = %5Bfragment.text%5D%0A    for child in fragment:%0A      body.append(lxml.html.tostring(child))%0A      body.append(child.tail)%0A    body = u''.join(body)%0A
 %0A   

allow to let <script> tag within deferJS (for editor syntax highlight for example), and remove it when collecting fragment