#!/usr/bin/env python3

# 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.

# Script to take raw sample output from Servo sampling profiler and
# output a [processed profile]. Based largely on [this script] and
# [this documentation].
#
# [processed profile]: https://github.com/firefox-devtools/profiler/blob/main/docs-developer/processed-profile-format.md
# [this script]: https://github.com/firefox-devtools/profiler/blob/main/src/profile-logic/import/linux-perf.js
# [this documentation]: https://github.com/firefox-devtools/profiler/blob/main/src/types/profile.js

from collections import defaultdict
import json
import sys


class StringTable:
    def __init__(self):
        self.table = {}
        self.idx = 0

    def get(self, s):
        assert s
        if s not in self.table:
            self.table[s] = self.idx
            self.idx += 1
        return self.table[s]

    def length(self):
        return len(list(self.table.keys()))

    def contents(self):
        return sorted(list(self.table.keys()), key=self.table.__getitem__)


with open(sys.argv[1]) as f:
    profile = json.load(f)
    rate = profile["rate"]
    samples = profile["data"]
    startTime = profile["start"]

frames = {}
stacks = {}

thread_data = defaultdict(list)
thread_order = {}
for sample in samples:
    if sample["name"]:
        name = sample["name"]
    else:
        name = "%s %d %d" % (sample["type"], sample["namespace"], sample["index"])
    thread_data[name].append((sample["time"], sample["frames"]))
    if name not in thread_order:
        thread_order[name] = (sample["namespace"], sample["index"])

tid = 0
threads = []
for name, raw_samples in sorted(iter(thread_data.items()), key=lambda x: thread_order[x[0]]):
    string_table = StringTable()
    tid += 1

    stackMap = {}
    stacks = []
    frameMap = {}
    frames = []

    samples = []

    for sample in raw_samples:
        prefix = None
        for frame in sample[1]:
            if not frame["name"]:
                continue
            if frame["name"] not in frameMap:
                frameMap[frame["name"]] = len(frames)
                frame_index = string_table.get(frame["name"])
                frames.append([frame_index])
            frame = frameMap[frame["name"]]

            stack_key = "%d,%d" % (frame, prefix) if prefix else str(frame)
            if stack_key not in stackMap:
                stackMap[stack_key] = len(stacks)
                stacks.append([frame, prefix])
            stack = stackMap[stack_key]
            prefix = stack
        samples.append([stack, sample[0]])

    threads.append(
        {
            "tid": tid,
            "name": name,
            "markers": {
                "schema": {
                    "name": 0,
                    "time": 1,
                    "data": 2,
                },
                "data": [],
            },
            "samples": {
                "schema": {
                    "stack": 0,
                    "time": 1,
                    "responsiveness": 2,
                    "rss": 2,
                    "uss": 4,
                    "frameNumber": 5,
                },
                "data": samples,
            },
            "frameTable": {
                "schema": {
                    "location": 0,
                    "implementation": 1,
                    "optimizations": 2,
                    "line": 3,
                    "category": 4,
                },
                "data": frames,
            },
            "stackTable": {
                "schema": {
                    "frame": 0,
                    "prefix": 1,
                },
                "data": stacks,
            },
            "stringTable": string_table.contents(),
        }
    )


output = {
    "meta": {
        "interval": rate,
        "processType": 0,
        "product": "Servo",
        "stackwalk": 1,
        "startTime": startTime,
        "version": 4,
        "presymbolicated": True,
    },
    "libs": [],
    "threads": threads,
}

print(json.dumps(output))