#!/usr/bin/env python

# Copyright 2018 The Servo Project Developers. See the COPYRIGHT
# file at the top-level directory of this distribution.
#
# Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
# http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
# <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
# option. This file may not be copied, modified, or distributed
# except according to those terms.

import json
import os
import sys
import tempfile
import webbrowser


def extract_memory_reports(lines):
    in_report = False
    report_lines = []
    times = []
    for line in lines:
        if line.startswith('Begin memory reports'):
            in_report = True
            report_lines += [[]]
            times += [line.strip().split()[-1]]
        elif line == 'End memory reports\n':
            in_report = False
        elif in_report:
            if line.startswith('|'):
                report_lines[-1].append(line.strip())
    return (report_lines, times)


def parse_memory_report(lines):
    reports = {}
    parents = []
    last_separator_index = None
    for line in lines:
        assert(line[0] == '|')
        line = line[1:]
        if not line:
            continue
        separator_index = line.index('--')
        if last_separator_index and separator_index <= last_separator_index:
            while parents and parents[-1][1] >= separator_index:
                parents.pop()

        amount, unit, _, name = line.split()

        dest_report = reports
        for (parent, index) in parents:
            dest_report = dest_report[parent]['children']
        dest_report[name] = {
            'amount': amount,
            'unit': unit,
            'children': {}
        }

        parents += [(name, separator_index)]
        last_separator_index = separator_index
    return reports


def transform_report_for_test(report):
    transformed = {}
    remaining = list(report.items())
    while remaining:
        (name, value) = remaining.pop()
        transformed[name] = '%s %s' % (value['amount'], value['unit'])
        remaining += map(lambda (k, v): (name + '/' + k, v), list(value['children'].items()))
    return transformed


def test_extract_memory_reports():
    input = ["Begin memory reports",
             "|",
             "  154.56 MiB -- explicit\n",
             "|     107.88 MiB -- system-heap-unclassified\n",
             "End memory reports\n"]
    expected = ([['|', '|     107.88 MiB -- system-heap-unclassified']], ['reports'])
    assert(extract_memory_reports(input) == expected)
    return 0


def test():
    input = '''|
|   23.89 MiB -- explicit
|      21.35 MiB -- jemalloc-heap-unclassified
|       2.54 MiB -- url(https://servo.org/)
|          2.16 MiB -- js
|             1.00 MiB -- gc-heap
|                0.77 MiB -- decommitted
|             1.00 MiB -- non-heap
|          0.27 MiB -- layout-thread
|             0.27 MiB -- stylist
|          0.12 MiB -- dom-tree
|
|   25.18 MiB -- jemalloc-heap-active'''

    expected = {
        'explicit': '23.89 MiB',
        'explicit/jemalloc-heap-unclassified': '21.35 MiB',
        'explicit/url(https://servo.org/)': '2.54 MiB',
        'explicit/url(https://servo.org/)/js': '2.16 MiB',
        'explicit/url(https://servo.org/)/js/gc-heap': '1.00 MiB',
        'explicit/url(https://servo.org/)/js/gc-heap/decommitted': '0.77 MiB',
        'explicit/url(https://servo.org/)/js/non-heap': '1.00 MiB',
        'explicit/url(https://servo.org/)/layout-thread': '0.27 MiB',
        'explicit/url(https://servo.org/)/layout-thread/stylist': '0.27 MiB',
        'explicit/url(https://servo.org/)/dom-tree': '0.12 MiB',
        'jemalloc-heap-active': '25.18 MiB',
    }
    report = parse_memory_report(input.split('\n'))
    transformed = transform_report_for_test(report)
    assert(sorted(transformed.keys()) == sorted(expected.keys()))
    for k, v in transformed.items():
        assert(v == expected[k])
    test_extract_memory_reports()
    return 0


def usage():
    print('%s --test - run automated tests' % sys.argv[0])
    print('%s file - extract all memory reports that are present in file' % sys.argv[0])
    return 1


if __name__ == "__main__":
    if len(sys.argv) == 1:
        sys.exit(usage())

    if sys.argv[1] == '--test':
        sys.exit(test())

    with open(sys.argv[1]) as f:
        lines = f.readlines()
    (reports, times) = extract_memory_reports(lines)
    json_reports = []
    for (report_lines, seconds) in zip(reports, times):
        report = parse_memory_report(report_lines)
        json_reports += [{'seconds': seconds, 'report': report}]
    with tempfile.NamedTemporaryFile(delete=False) as output:
        thisdir = os.path.dirname(os.path.abspath(__file__))
        with open(os.path.join(thisdir, 'memory_chart.html')) as template:
            content = template.read()
            output.write(content.replace('[/* json data */]', json.dumps(json_reports)))
            webbrowser.open_new_tab('file://' + output.name)