#!/usr/bin/env python3

import argparse
import json
import re
import sys
import yaml

def log(message):
    sys.stderr.write(message + "\n")

def is_metaaddress(value):
    return (isinstance(value, str)
            and value.startswith("0x")
            and (":Code_" in value
                 or ":Generic" in value))

class MetaAddressRemapper:
    def __init__(self):
        self.addresses = set()

    def handle(self, value):
        if is_metaaddress(value):
            self.addresses.add(value)

    def collect(self, value):
        if isinstance(value, dict):
            for k, v in value.items():
                self.handle(k)
                self.collect(v)

        elif isinstance(value, list):
            for v in value:
                self.collect(v)

        else:
            self.handle(value)

    def apply_replacement(self, value):
        return (self.replacements[value]
                if (isinstance(value, str)
                    and value in self.replacements)
                else value)

    def replace(self, value):
        if isinstance(value, dict):
            new_values = {self.apply_replacement(k) :
                          self.apply_replacement(v)
                          for k, v
                          in value.items()}
            value.clear()
            value.update(new_values)

            for k, v in value.items():
                self.replace(v)

        elif isinstance(value, list):
            new_values = [self.apply_replacement(v) for v in value]
            value.clear()
            for new_value in new_values:
                value.append(new_value)

            for v in value:
                self.replace(v)

    def rewrite(self, value):
        self.replacements = {v: str(i)
                             for i, v
                             in enumerate(sorted(self.addresses))}
        self.replace(value)

        return value

class SafeLoaderIgnoreUnknown(yaml.SafeLoader):
    def ignore_unknown(self, node):
        return self.construct_mapping(node)

SafeLoaderIgnoreUnknown.add_constructor(None,
                                        SafeLoaderIgnoreUnknown.ignore_unknown)

def remap_metaaddress(model):
    mar = MetaAddressRemapper()
    mar.collect(model)
    return mar.rewrite(model)

def main():
    parser = argparse.ArgumentParser(description="Extract and process rev.ng model.")
    parser.add_argument("--json",
                        action="store_true",
                        help="Dump as JSON.")
    parser.add_argument("--remap",
                        action="store_true",
                        help="Remap MetaAddresses. Implies --json.")
    args = parser.parse_args()

    if args.remap:
        args.json = True

    prefix = None

    # Find revng.model named metadata
    for line in sys.stdin:
        match = re.match("!revng.model = !{!([0-9]*)}", line)
        if match:
            prefix = "!" + match.groups(1)[0] + " = !{!\""
            break

    # Early exit if not found
    if not prefix:
        log("Couldn't find revng.model metadata")
        return 1

    # Look for associated named metadata5~
    for line in sys.stdin:
        if prefix and line.startswith(prefix):
            text_model = line[len(prefix):-3]

            # Unescape the string
            for escaped in set(re.findall(r"\\[0-9a-fA-F]{2}", text_model)):
                replacement = bytearray.fromhex(escaped[1:]).decode()
                text_model = text_model.replace(escaped, replacement)

            if not args.json:
                print(text_model)
                return 0

            # Decode YAML
            model = yaml.load(text_model, Loader=SafeLoaderIgnoreUnknown)

            # Remap MetaAddress
            if args.remap:
                model = remap_metaaddress(model)

            # Dump as JSON
            print(json.dumps(model,
                             indent=2,
                             sort_keys=True,
                             check_circular=False))

            # Consume all remaining input
            sys.stdin.read()

            return 0

    log("Couldn't find revng.model metadata")
    return 1

if __name__ == "__main__":
    sys.exit(main())