# Copyright 2010 New Relic, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import re

_head_re = re.compile(b'<head[^>]*>', re.IGNORECASE)

_xua_meta_re = re.compile(b"""<\s*meta[^>]+http-equiv\s*=\s*['"]"""
    b"""x-ua-compatible['"][^>]*>""", re.IGNORECASE)

_charset_meta_re = re.compile(b"""<\s*meta[^>]+charset\s*=[^>]*>""",
    re.IGNORECASE)

_attachment_meta_re = re.compile(b"""<\s*meta[^>]+http-equiv\s*=\s*['"]"""
    b"""content-disposition['"][^>]*content\s*=\s*(?P<quote>['"])"""
    b"""\s*attachment(\s*;[^>]*)?(?P=quote)[^>]*>""",
    re.IGNORECASE)

_body_re = re.compile(b'<body[^>]*>', re.IGNORECASE)

def insert_html_snippet(data, html_to_be_inserted, search_limit=64*1024):
    # First determine if we have a body tag. If we don't we
    # always give up even though strictly speaking we may not
    # actually need it to exist. This is to ensure that we have
    # all the HTML needed to perform the match correctly. In
    # doing this initial search, we only do up to the specified
    # search limit.

    body = _body_re.search(data[:search_limit])

    if not body:
        return data if len(data) > search_limit else None

    # We are definitely doing to insert something now, so
    # generate the text to be inserted. Bail out if is empty but
    # return the data to indicate we did insert something.

    text = html_to_be_inserted()

    if not text:
        return data

    # We are now going to split out everything before the body
    # to limit searches when applying the regex patterns. We
    # could have just used up to the search limit, but the
    # patterns which follow seem to be expensive enough that a
    # short a string as possible helps.

    tail, data = data[body.start():], data[:body.start()]

    def insert_at_index(index):
       return b''.join((data[:index], text, data[index:], tail))

    # Search for instance of a content disposition meta tag
    # indicating that the response is actually being served up
    # as an attachment and would be saved as a file and not
    # actually interpreted by a browser.

    if _attachment_meta_re.search(data):
        return data + tail

    # Search for instances of X-UA or charset meta tags. We will
    # use whichever is the last to appear in the data.

    xua_meta = _xua_meta_re.search(data)
    charset_meta = _charset_meta_re.search(data)

    index = max(xua_meta and xua_meta.end() or 0,
            charset_meta and charset_meta.end() or 0)

    if index:
       return insert_at_index(index)

    # Next try for the start of head section.

    head = _head_re.search(data)

    if head:
       return insert_at_index(head.end())

    # Finally if no joy, insert before the start of the body.

    return insert_at_index(body.start())

def verify_body_exists(data):
    return _body_re.search(data)
