#!/usr/bin/env python3

# Written by: Chaitanya S Lakkundi (chaitanya.lakkundi@gmail.com)

import re
from anytree import Node as NodeParent, PreOrderIter
from anytree.util import rightsibling
from html import escape, unescape
import json

INDENT = 2

class Node(NodeParent):
    tag_prefix = {"va": "v.", "sh": "sh.", "adhyaya": "ady.", "su": "su.", "cm": "c."}
    prop_order = ['id', 'rel_id']
    INDENT = 2

    def __init__(self, raw_data, parent=None):
        super().__init__(raw_data, parent)
        if not raw_data.strip():
            raise Exception("Empty data. Cannot create node")
        self.raw_data = raw_data
        self.tag_name, core, self.content = self.parse_ikml_node(raw_data)
        # first see if key="val" this pattern exists, only then after removing all of them split by space and check
        # with quotes, wq and without quotes, woq
        props_wq = re.findall(r"([A-z]+)=[\"\'](.*?)[\"\']", core)
        props_wq = {k:v for k, v in props_wq}
        core_woq = re.sub(r"[\"\'](.*?)[\"\']", "<!!>", core)
        props_woq = re.findall(r"([A-z]+)=(.*?)[ \]]", core_woq)
        props_woq = {k:v for k, v in props_woq if v != "<!!>"}

        # print(core_woq, props_wq, props_woq)
        self.props = dict()
        self.props.update(props_wq)
        self.props.update(props_woq)

    def keys(self):
        ret = list(self.props.keys())
        for n in self.children:
            if n.tag_name and n.tag_name[0] == ".":
                ret.append(n.tag_name)
        return ret
    
    def get(self, key, default=None):
        try:
            return self[key]
        except:
            return default

    def __getitem__(self, key):
        # remove dot in the beginning to make it searcheable in both attribute formats
        key = key.replace(".", "")
        try:
            return self.props[key]
        except:
            pass

        for n in self.children:
            if n.tag_name == "."+key:
                return n.content

        raise KeyError(f"'{key}' attribute doesn't exist in node {self.raw_data}")

    def __setitem__(self, key, value):
        if key.startswith("."):
            Node(f"[{key}] {value}", parent=self)
        else:
            self.props[key] = value

    def __str__(self, quoted_attr=True):
        buf = [self.tag_name]
        # for pn in self.valid_props:
        keys = self.props.keys() - self.prop_order
        for pn in self.prop_order + sorted(keys):
            try:
                if quoted_attr:
                    # ESCAPING XML https://stackoverflow.com/a/46637835
                    # WARNING: keep double quotes only to support apostrophes in attrib value. eg. Opponent's View
                    buf.append(f"{pn}=\"{self.props[pn]}\"")
                else:
                    buf.append(f"{pn}={self.props[pn]}")
            except:
                pass
        t = "[" + " ".join(buf) + "]"
        if self.content:
            t = t + " " + self.content
        return t

    @property
    def node_children(self):
        for child in self.children:
            if not child.is_attribute:
                yield child

    @property
    def is_attribute(self):
        return self.tag_name.startswith(".")

    def generate_id(self):
        # means, prop id doesn't exist, then borrow from parent
        if self.parent != None and self.parent.props.get("id", None):
            if "rel_id" in self.props:
                self.props["id"] = self.parent.props["id"] + "." + self.tag_prefix.get(self.tag_name, "") + str(self.props["rel_id"])
            else:
                self["id"] = self.parent["id"]
        # if no parent, check for rel_id. make it the id.
        elif "rel_id" in self.props:
            self["id"] = str(self["rel_id"])

    def generate_ids(self):
        # TODO: find a better way to express the following
        try:
            self["id"] = str(self["root_prefix"]) + "." + str(self["rel_id"])
        except:
            try:
                self["id"] = str(self["rel_id"])
            except:
                pass

        for cn in self.children:
            for node in PreOrderIter(cn):
                node.generate_id()

    def generate_relids(self, root_rel_id=None):
        # Logic
        # For every prefix, keep a separate rel_id count. Traverse sibling nodes and generate rel_ids
        # Now do the same for children of every sibling node
        # 0 is null tag
        # relcount is reset for every child. but kept same for siblings
        relcounts = {t:0 for t,p in self.tag_prefix.items()}
        relcounts[0] = 0
        # ikml_out.append(str(root))
        try:
            sibl = self.children[0]
        except:
            return
        while sibl:
            if sibl.is_attribute:
                sibl = rightsibling(sibl)
                continue
            try:
                rel_id = sibl["rel_id"]
            except KeyError:
                rel_id = None
            # if rel_id is an integer, we can overwrite
            if rel_id is None or re.search(r"^[0-9]+$", str(rel_id)):
                if sibl.tag_name in self.tag_prefix.keys():
                    key = sibl.tag_name
                else:
                    key = 0
                relcounts[key] += 1
                sibl["rel_id"] = relcounts[key]
            else:
                # means, explicitly defined rel_id; skip
                pass
            sibl.generate_relids()
            sibl = rightsibling(sibl)

    def put_attrs_inside(self, dotted=False):
        # complies to xml strictly. means, attributes are not child nodes anymore. they become properties within opening node itself.
        for child in self.children:
            if child.is_attribute:
                if dotted:
                    child.parent[child.tag_name] = self.parse_ikml_node(str(child))[2]
                else:
                    child.parent[child.tag_name[1:]] = self.parse_ikml_node(str(child))[2]
                child.parent = None
                continue
            child.put_attrs_inside(dotted)

    def to_xml(self, quoted_attr=True):
        root = self.__class__("[dummy_root]")
        root.children = [self]
        return "\n".join(self.__class__.tree_as_xml_list(root, quoted_attr))

    def to_dict(self, recurse_on_children=True):
        dict_out = {"tag_name": self.tag_name, "content": self.content, "children": []}
        for k, v in self.props.items():
            dict_out[k] = v

        for child in self.children:
            if child.is_attribute:
                # do not remove the dot in tag_name
                # attr_name = child.tag_name[1:]
                dict_out[child.tag_name] = child.content
                continue
            
        for child in self.children:
            if recurse_on_children and not child.is_attribute:
                dict_out["children"].append(child.to_dict(recurse_on_children=recurse_on_children))
        
        return dict_out

    def to_json(self, recurse_on_children=True):
        d = self.to_dict(recurse_on_children=recurse_on_children)
        return json.dumps(d, ensure_ascii=False, indent=self.INDENT)
    
    def to_txt(self):
        d = self.tree_as_list(self.root, quoted_attr=False)
        return "\n".join(d)

    @classmethod
    def tree_as_list(cls, root, dummy_root=False, quoted_attr=True):
        if dummy_root:
            ikml_out = []
            adjusted = 1
        else:
            ikml_out = [cls.__str__(root, quoted_attr)]
            adjusted = 0
        for cn in root.children:
            for node in PreOrderIter(cn):
                ikml_out.append((node.depth-adjusted)*cls.INDENT*" " + cls.__str__(node, quoted_attr))
        return ikml_out

    @classmethod
    def tree_as_xml_list(cls, root, quoted_attr=True):
        xmlout = []
        # only top level root children are considered
        for child in root.children:
            xmlout.append(cls.open_tag(cls.__str__(child, quoted_attr)))
            xmlout.extend(cls.tree_as_xml_list(child, quoted_attr) or [])
            xmlout.append(cls.close_tag(cls.__str__(child, quoted_attr)))
        return xmlout

    @staticmethod
    def parse_ikml_node(tagline):
        try:
            content = unescape(tagline.split("]", 1)[1].strip())
            core = re.findall(r"(\[.*?\])", tagline)[0].strip()
            tag_name = re.search(r"\[[ ]*([.A-z]+).*?\]", core).groups()[0].strip()
            return (tag_name, core, content)
        except:
            return ("", "[]", "")

    @classmethod
    def open_tag(cls, tagline):
        return tagline.strip().replace("[", "<").replace("]", ">")

    @classmethod
    def close_tag(cls, tagline):
        if not tagline.strip():
            return "</>"
        if tagline.strip()[0] != "[":
            return "</"+tagline+">"
        tag_name, core, content = cls.parse_ikml_node(tagline)
        return "</"+tag_name+">"

def fix_root_prefix(root, root_prefix):
    # change is inplace for anytree Node
    if root_prefix:
        try:
            root["id"] = root_prefix + "." + root["rel_id"]
        except:
            root["id"] = root_prefix
        root["root_prefix"] = root_prefix

def ikml_to_anytree(ikml_data, root_prefix=None):
    root = Node(f"[root]")
    parents_stack = []
    prev_nd = root
    lvl = -1

    if isinstance(ikml_data, str):
        ikml_data = ikml_data.splitlines()

    for line in ikml_data:
        if not line.split("#")[0].strip():
            continue
        tindent = line.find("[")
        cur_lvl = int(tindent / INDENT)
        if cur_lvl > lvl:
            parents_stack.append(prev_nd)
        elif cur_lvl < lvl:
            try:
                while lvl != cur_lvl:
                    parents_stack.pop()
                    lvl -= 1
            except Exception as e:
                msg = "error in ikml_to_anytree. this or prev tag: " + line
                print(msg)
                break
        try:
            nd = Node(line.strip(), parent=parents_stack[-1])
            prev_nd = nd
        except Exception as e:
            msg = "error in ikml_to_anytree. this or prev tag: " + line
            print(msg)
            # raise e
            break
        lvl = cur_lvl
    
    fix_root_prefix(root, root_prefix)

    return root

def dict_to_anytree(ikml_dicts, root_prefix=None):
    # ikml_dicts is a list of dict items
    if isinstance(ikml_dicts, dict):
        ikml_dicts = [ikml_dicts]
    
    root = Node(f"[root]")

    # set the root props
    # root content must always be empty
    if ikml_dicts[0]["tag_name"] == "root":
        props = set(ikml_dicts[0].keys()) - set(('tag_name', 'children', 'content'))
        for p in props:
            root[p] = ikml_dicts[0][p]
        iterate_over_children = ikml_dicts[0]['children']
    else:
        iterate_over_children = ikml_dicts

    def _dict_to_anytree(parent, children):
        if children is not None:
            for child in children:
                props = set(child.keys()) - set(('tag_name', 'children', 'content'))
                n = Node(f"[{child['tag_name']}]", parent=parent)
                n.content = child.get('content', '')
                for p in props:
                    if p.startswith("."):
                        Node(f"[{p}] {child[p]}")
                    else:
                        n[p] = child[p]
                _dict_to_anytree(child, child.get('children'))

    _dict_to_anytree(root, iterate_over_children)

    fix_root_prefix(root, root_prefix)

    return root


if __name__ == "__main__":
    from sys import argv
    from pprint import pprint
    try:
        input_file = argv[1]
        ikml_data = open(input_file).read().splitlines()
        root = ikml_to_anytree(ikml_data)
        pprint(Node.tree_as_list(root))
    except IndexError:
        print(f"Usage: {argv[0]} input_ikml_file.txt")
