# Node data classes for parsing and synthesis

# The parser emits a tree of Node subclasses
# and further synthesis steps operate on such trees

import os

from .structures import DataClass


class Node(DataClass):
    pass


# Nodes generated by the parser


class RootNode(Node):
    _attributes = ["modules", "config"]


class ModuleNode(Node):
    _attributes = [
        "name",
        "ports",
        "body",
        "symbols",
        "nets",
        "gates",
        "continuous_assignments",
        "net_counter",
        "inst_counter",
    ]
    _defaults = {
        "body": None,
        "symbols": None,
        "nets": None,
        "gates": None,
        "continuous_assignments": None,
        "net_counter": None,
        "inst_counter": None,
    }


class PortNode(Node):
    _attributes = ["name", "direction", "signed", "data_type", "index", "term"]
    _defaults = {"signed": None, "data_type": None, "index": None, "term": None}


class NumericNode(Node):
    _attributes = ["value", "signed", "size"]


class IdentifierNode(Node):
    _attributes = ["name"]


class HierarchicalIdentifierNode(Node):
    _attributes = ["parent", "name"]


class ConcatenationNode(Node):
    _attributes = ["expressions"]


class ReplicationNode(Node):
    _attributes = ["repeat", "expressions"]


class UnaryExpressionNode(Node):
    _attributes = ["operator", "argument"]


class BinaryExpressionNode(Node):
    _attributes = ["operator", "arguments"]


class ConditionalExpressionNode(Node):
    _attributes = ["condition", "branch0", "branch1"]


class IndexExpressionNode(Node):
    _attributes = ["identifier", "index"]


class RangeNode(Node):
    _attributes = ["operator", "indices"]


class DeclarationNode(Node):
    _attributes = ["name", "signed", "data_type", "index", "assignment"]


class EventExpressionNode(Node):
    _attributes = ["edge", "expression"]


class AlwaysBlockNode(Node):
    _attributes = ["block_type", "sensitivity_list", "statement"]


class NullStatementNode(Node):
    _attributes = []


class AssignmentStatementNode(Node):
    _attributes = ["blocking", "lhs", "rhs"]


class ConditionalStatementNode(Node):
    _attributes = ["condition", "branch0", "branch1"]


class CaseStatementNode(Node):
    _attributes = ["expression", "cases"]


class CaseItemNode(Node):
    _attributes = ["matches", "branch"]


class ForeverStatementNode(Node):
    _attributes = ["body"]


class RepeatStatementNode(Node):
    _attributes = ["count", "body"]


class WhileStatementNode(Node):
    _attributes = ["condition", "body"]


class ForStatementNode(Node):
    _attributes = ["init", "condition", "step", "body"]


class SequentialBlockNode(Node):
    _attributes = ["statements"]


class ContinuousAssignmentNode(Node):
    _attributes = ["lhs", "rhs"]


class GateInstantiationNode(Node):
    _attributes = ["gate", "input_terminals", "output_terminal"]


class ModuleInstantiationNode(Node):
    _attributes = ["module_name", "instance_name", "instance_index", "port_connections"]


class PortConnectionNode(Node):
    _attributes = ["port_name", "value"]


# Nodes generated by later steps in synthesis


class NetNode(Node):
    _attributes = ["name", "terms", "driver"]
    _defaults = {"terms": None, "driver": None}


class BitNode(Node):
    _attributes = ["value"]


class DontCareNode(Node):
    _attributes = []


class BusNode(Node):
    _attributes = ["signed", "bits"]


class ScopedBusNode(Node):
    _attributes = ["gates", "signed", "bits"]


class GateNode(Node):
    _attributes = ["name", "instance", "terminals", "terms"]
    _defaults = {"terminals": None, "terms": None}


class SymbolNode(Node):
    _attributes = ["name", "signed", "data_type", "index"]


class NetlistNode(Node):
    _attributes = ["name", "num_terms", "nets", "ports", "instances", "config"]


# Utility functions and classes for nodes


def node_type(node):
    """get snake_case version of node class basename (e.g. ConditionalExpressionNode => conditional_expression)"""
    try:
        return getattr(node, "_type")
    except AttributeError:
        assert isinstance(node, Node)
        type_str = "".join(
            ("_" + c.lower() if c.isupper() else c) for c in node.__class__.__name__
        )
        assert type_str.startswith("_")
        assert type_str.endswith("_node")
        assert type_str.islower()
        type_str = type_str[1:-5]
        node.__class__._type = type_str
        return type_str


class NodeProcessor:
    """base class for a synthesis step recursively processing a syntax tree"""

    def __init__(self):
        pass

    def process(self, node):
        """relay processing to a specialized function process_* depending on node type"""
        method_name = "process_" + node_type(node)
        try:
            method = getattr(self, method_name)
        except AttributeError:
            raise NotImplementedError(
                f"Node processor {self.__class__.__name__} not implemented for node type {node.__class__.__name__}"
            ) from None
        return method(node)

    def process_obj(self, obj):
        """recursively process lists, tuples, dicts containing nodes"""
        if obj is None or type(obj) in (bool, int, float, range, str):
            return obj
        elif isinstance(obj, Node):
            return self.process(obj)
        elif type(obj) is list:
            return [self.process_obj(e) for e in obj]
        elif type(obj) is tuple:
            return tuple(self.process_obj(e) for e in obj)
        elif type(obj) is dict:
            return {dk: self.process_obj(dv) for dk, dv in obj.items()}
        elif isinstance(obj, getattr(os, "PathLike", type(None))):
            return obj
        elif isinstance(obj, type):
            return obj
        else:
            raise RuntimeError(
                f"Node processor {self.__class__.__name__} has no pass-through implementation for {obj}"
            )

    def pass_through(self, node):
        """don't change the current node but keep processing its children"""
        props = {}
        for k, v in node.__dict__.items():
            if not k.startswith("_"):
                props[k] = self.process_obj(v)
        return node.__class__(**props)

    def debug_props(self, node):
        """dump properties for current node, useful as a first step of implementing a process_* function"""
        props = [e for e in dir(node) if not e.startswith("_")]
        props_str = ", ".join(props)
        print(f"{node.__class__.__name__}: {props_str}")
        return node

    def create_net(self):
        try:
            nn = NetNode(name=f"$net_{self.net_counter}")
            self.net_counter += 1
            self.nets.append(nn)
            return nn
        except AttributeError:
            raise NotImplementedError(
                f"Node processor {self.__class__.__name__} doesn't keep track of nets"
            ) from None

    def create_instance_name(self):
        try:
            name = f"$inst_{self.inst_counter}"
            self.inst_counter += 1
            return name
        except AttributeError:
            raise NotImplementedError(
                f"Node processor {self.__class__.__name__} doesn't keep track of instances"
            ) from None

    def create_gate(self, name, outputs=None, **terminals):
        instance = self.create_instance_name()
        if outputs is None:
            output_nets = None
            terminals = terminals
        elif type(outputs) is str:
            output_nets = self.create_net()
            terminals = terminals | {outputs: output_nets}
        else:
            output_nets = tuple(self.create_net() for s in outputs)
            terminals = terminals | dict(zip(outputs, output_nets))
        nn = GateNode(name=name, instance=instance, terminals=terminals)
        try:
            self.gates.append(nn)
        except AttributeError:
            raise NotImplementedError(
                f"Node processor {self.__class__.__name__} doesn't keep track of gates"
            ) from None
        return output_nets


def assert_net_or_bit(*args):
    for arg in args:
        if isinstance(arg, NetNode):
            continue
        elif isinstance(arg, BitNode) and arg.value in (0, 1):
            continue
        elif isinstance(arg, DontCareNode):
            continue
        else:
            raise SyntaxError(f"Expected net or constant bit, found {arg}")


def dump_text(text, end="\n", file=None):
    """helper function for dump_tree"""
    if file is None:
        print(text, end=end)
    else:
        print(text, end=end, file=file)


def dump_tree(obj, file=None, indent=0, prefix="- ", skip_indent=False):
    """prints a node and its subtree in human-readable format"""
    sp = ("" if skip_indent else "    " * indent) + prefix
    if isinstance(obj, Node):
        filtered = []
        for k, v in sorted(obj.__dict__.items()):
            assert type(k) is str
            if k.startswith("_"):
                continue
            if k in obj._defaults:
                if obj._defaults[k] == v:
                    continue
            filtered.append((k, v))
        obj_type = node_type(obj).upper()
        if len(filtered) == 1:
            k, v = filtered[0]
            dump_tree(
                v,
                file=file,
                indent=indent,
                prefix=f"{sp}{obj_type} {k}: ",
                skip_indent=True,
            )
        else:
            dump_text(f"{sp}{obj_type}", file=file)
            for k, v in filtered:
                dump_tree(v, file=file, indent=indent + 1, prefix=f"- {k}: ")
    elif obj is None:
        dump_text(f"{sp}(None)", file=file)
    elif type(obj) in (bool, int, range, str):
        dump_text(f"{sp}{obj}", file=file)
    elif type(obj) is list:
        dump_text(f"{sp}list", file=file)
        for e in obj:
            dump_tree(e, file=file, indent=indent + 1)
    elif type(obj) is tuple:
        dump_text(f"{sp}tuple", file=file)
        for e in obj:
            dump_tree(e, file=file, indent=indent + 1)
    elif type(obj) is set:
        dump_text(f"{sp}set", file=file)
        for e in sorted(obj):
            dump_tree(e, file=file, indent=indent + 1)
    elif type(obj) is dict:
        dump_text(f"{sp}dict", file=file)
        for k, v in sorted(obj.items()):
            assert type(k) is str
            dump_tree(v, file=file, indent=indent + 1, prefix=f"- {k}: ")
    else:
        dump_text(f"{sp}{obj}", file=file)


def net_sort_key(name):
    """key for sorting net/instance names naturally, e.g. $net_9, $net_10, $net_11"""
    # "a123bc4" => ("a", 123, "bc", 4)
    key = [""]
    last_digit = False
    for c in name:
        cur_digit = c.isdigit()
        if cur_digit == last_digit:
            key[-1] += c
        elif last_digit:
            key[-1] = int(key[-1])
            key.append(c)
        else:
            key.append(c)
        last_digit = cur_digit
    if last_digit:
        key[-1] = int(key[-1])
    return tuple(key)
