from django import template
import re
from django.contrib.flatpages.models import FlatPage
register = template.Library()

def do_flatpage_root(parser, token):
    """
    This tag creates a queryset containing all the top-level flatpages,
    adding it to the current context with the given name.
    This does not include "/", if there is a flatpage with that url.
    
    For example::
    
        {% flatpage_root as root %}
    
    adds a queryset of flatpages whose url can be matched with ^/[^/]*/$ to
    the current context, named root. It can then be used::
    
        {% for page in root %}
        
    and so forth.
    """
    bits = token.contents.split()
    if len(bits) != 3 or bits[1] != "as":
        raise template.TemplateSyntaxError("%r expected format is 'flatpage_root as name'" % bits[0])
    return FlatpageRoot(bits[2])

class FlatpageRoot(template.Node):
    # this regex is the definition of a top-level flatpage
    # override in a subclass if you have a different idea!
    root_regex = r'^/[^/]*/$'
    def __init__(self, context_name):
        self.context_name = context_name
    def render(self, context):
        context[self.context_name] = FlatPage.objects.filter(url__regex=self.root_regex)
        return ""

register.tag('flatpage_root', do_flatpage_root)

def do_flatpage_children(parser, token):
    """
    This tag creates a queryset containing all flatpages below a given root
    url, adding it to the current context with the given name.
    
    A "root url" is any url matched by ^/[^/]*/$.
    
    Typically, this would be used to get a list of children of a parent url::
    
        {% flatpage_children "/about/" as children}
    
    This adds a queryset of all deeper urls, such as "/about/author/" and
    "/about/pants/" to the current context, named children.
    
    If the url given is "/", the tag doesn't do anything.
    
    If the url given is not a root url, for example "/about/pants/", only the
    root url portion of the argument is used. Thus, "/about/" and
    "/about/pants/" will return the same result, assuming both those flatpages
    exist.
    
    If the first argument is not in quotes, it is assumed to be in the current
    context::

        {% flatpage_children flatpage.url as children %}
    
    This adds a queryset of all urls deeper than the root url of the current
    flatpage.
    """
    bits = token.contents.split()
    if len(bits) != 4 or bits[2] != "as":
        raise template.TemplateSyntaxError("%r expected format is 'flatpage_children URL as name'" % bits[0])
    return FlatpageChildren(bits[1], bits[3])

class FlatpageChildren(template.Node):
    # this regex is the definition of a top-level flatpage
    # override in a subclass if you have a different idea!
    # note the difference from root_regex in FlatpageRoot:
    # the string is allowed to continue past the first slash...
    root_regex = r'^(/[^/]*?/)'
    def __init__(self, url, context_name):
        self.url = url
        self.context_name = context_name
    def render(self, context):
        if self.url[0] in ('"',"'"):
            if self.url[0] == self.url[-1] and len(self.url) > 3:
                urlcontent = self.url
            else:
                return ""
        else:
            urlcontent = template.Variable(self.url).resolve(context)
        m = re.search(self.root_regex, urlcontent)
        if not m:
            return ""
        root = m.group(1)
                
        context[self.context_name] = FlatPage.objects.filter(
            url__gt=root,
            url__startswith=root
        )
        return ""

register.tag('flatpage_children', do_flatpage_children)