#!/usr/bin/env python3
import argparse
import os
import tempfile
import subprocess
import json
from patracetools import headerparser


def patch(trace_file, header_file):
    with open(header_file) as hf:
        header = json.load(hf)

    print("Patching {trace}...".format(trace=trace_file))
    headerparser.write_json_header(trace_file, header)


def mergelist(original, update):
    for i, obj in enumerate(original):
        if isinstance(obj, dict):
            merge(obj, update[i])


def merge(original, update):
    """ Merges two python objects """
    for key, value in original.items():
        if key not in update:
            update[key] = value
        elif isinstance(value, dict):
            merge(value, update[key])
            return update
        elif isinstance(value, list):
            mergelist(value, update[key])

    return update


def edit(json_string):
    pretty_printed_json = json.dumps(json_string, indent=2, sort_keys=True)
    tmp_header = tempfile.NamedTemporaryFile('wt', delete=False)
    tmp_header.write(pretty_printed_json)
    tmp_header.close()

    editor = os.environ.get('EDITOR', 'vi')
    p = subprocess.Popen([editor, tmp_header.name])
    p.communicate()

    # Un-prettyprint the json data
    with open(tmp_header.name) as tmpfile:
        json_string = json.loads(tmpfile.read())

    with open(tmp_header.name, 'w') as tmpfile:
        tmpfile.write(json.dumps(json_string, separators=(',', ':')))

    return tmpfile.name


def find_thread(threads, id):
    for thread in threads:
        if thread["id"] == id:
            return thread


def main():
    parser = argparse.ArgumentParser(description='Edit the header of a trace, in-place.')
    parser.add_argument('file', help='Path to the .pat trace file')
    parser.add_argument('--merge', help='If specified, a path to a json file that will be merged into existing header')
    parser.add_argument('--mergeattrpatch', help='Merge patch file generated by attribute_result_to_header.py')
    args = parser.parse_args()

    parsed = headerparser.read_json_header(args.file)

    if args.merge:
        with open(args.merge) as f:
            data = json.load(f)

        merged = merge(parsed, data)

        jsonstr = json.dumps(merged, separators=(',', ':'))
        tmpfile = tempfile.NamedTemporaryFile(delete=False)
        tmpfile.write(jsonstr)
        tmpfile.close()

        file_name_edited = tmpfile.name
    elif args.mergeattrpatch:
        with open(args.mergeattrpatch) as f:
            data = json.load(f)

        # Remove old attributes:
        for thread in parsed["threads"]:
            if "attributes" in thread:
                del thread["attributes"]

        for i, thread in enumerate(data["threads"]):
            if "attributes" in thread:
                orig_thread = find_thread(parsed["threads"], i)
                orig_thread["attributes"] = thread["attributes"]

        jsonstr = json.dumps(parsed, separators=(',', ':'))
        tmpfile = tempfile.NamedTemporaryFile(delete=False)
        tmpfile.write(jsonstr)
        tmpfile.close()

        file_name_edited = tmpfile.name
    else:
        # Use text editor
        file_name_edited = edit(parsed)

    patch(args.file, file_name_edited)

    os.unlink(file_name_edited)


if __name__ == '__main__':
    main()
