
import axiomic.protos as protos

from google.protobuf import descriptor
import json


from rich.console import Console
from rich.table import Table
from rich.tree import Tree
from rich.panel import Panel
from rich import box
import dataclasses

class Graph:
    def __init__(self, graph: protos.axiomic.AxiomicGraph):
        self.graph = graph
        self._nodes = {}

        for node in graph.nodes:
            self._nodes[node.name] = node

    def get_node(self, name: str) -> protos.axiomic.AxiomicNode:
        return self._nodes[name]

    def has_node(self, name: str) -> bool:
        return name in self._nodes

    def ref_to_node(self, ref: protos.axiomic.NodeRef) -> protos.axiomic.AxiomicNode:
        return self.get_node(ref.name)

    def flatten_node(self, node: protos.axiomic.AxiomicNode) -> list:
        return reversed_breadth_first_flatten(self, node)

    def print_tree(self, node: protos.axiomic.AxiomicNode):
        _print_tree(self, node)


def get_subnode(node: protos.axiomic.AxiomicNode) -> protos.axiomic.AxiomicNode:
    field_name = node.WhichOneof('node_type')
    sub_node = getattr(node, field_name)
    return field_name, sub_node


def is_axiomic_node(thing):
    return str(type(thing)) == str(protos.axiomic.NodeRef)


def list_format_node_childern_extras(graph: Graph, fn: protos.axiomic.FormatNode) -> list:
    l = []
    for format_input in fn.inputs:
        child = graph.ref_to_node(format_input.value)
        l.append((f'{{{format_input.name}}}', child))
    return l

def list_history_complete_node_childern_extras(graph: Graph, hn: protos.axiomic.HistoryCompleteNode) -> list:
    l = []
    for i, history_pair in enumerate(hn.chat_history):
        l.append((f'user_message_{i}', graph.ref_to_node(history_pair.user_message)))
        l.append((f'assistant_message_{i}', graph.ref_to_node(history_pair.assistant_message)))
    return l

SPECIAL_CASE_LIST = {
    'format_node': list_format_node_childern_extras,
    'history_complete_node': list_history_complete_node_childern_extras,
}

def list_childern(graph: Graph, node: protos.axiomic.AxiomicNode) -> list:
    sub_node_name, sub_node = get_subnode(node)
    descriptor = sub_node.DESCRIPTOR
    childern = []
    for field in descriptor.fields:
        field_name = field.name
        if is_axiomic_node(getattr(sub_node, field_name)):
            childern.append((field_name, graph.ref_to_node(getattr(sub_node, field_name))))
    
    if sub_node_name in SPECIAL_CASE_LIST:
        childern += SPECIAL_CASE_LIST[sub_node_name](graph, sub_node)
    return childern


def is_autogenerated_name(name: str) -> bool:
    return name.startswith('_')


def list_childern_named_only(graph: Graph, node: protos.axiomic.AxiomicNode) -> list:
    childern = list_childern(graph, node)

    # if a child has an autogenerated name, recurse to get another name for it.
    for i, (child_name, child_node) in enumerate(childern):
        ...


def str_parent(parent_name, childern):
    childern_str = ''
    for child in childern:
        childern_str += f'  {child[0]}: {child[1].name}'
    childern_str = childern_str.strip()
    return (parent_name, childern_str)


def abbrivate_str(s: str, max_len=80) -> str:
    if len(s) == 0:
        return '(EMPTY STRING)'
    s = s.strip().replace('\n', ' ').strip()
    if len(s) > max_len:
        return s[:max_len] + '...'
    return s


def is_leaf(graph: Graph, node: protos.axiomic.AxiomicNode) -> bool:
    return len(list_childern(graph, node)) == 0


def str_leaf(node: protos.axiomic.AxiomicNode) -> str:
    field_name = node.WhichOneof('node_type')

    # print(node)
    if field_name == 'immediate_node':
        abbrivated_val = abbrivate_str(node.immediate_node.value)
        name = node.name
        return name, abbrivated_val
    return name, 'LEAF'


def dedup_childern_list(childern: list) -> list:
    childern_set = set()
    deduped_childern = []
    for child in childern:
        if child.name not in childern_set:
            deduped_childern.append(child)
            childern_set.add(child.name)
    return deduped_childern


def _reversed_breadth_first_flatten(graph: Graph, node: protos.axiomic.AxiomicNode) -> list:
    childern = list_childern(graph, node)
    ans = []
    for child_name, child_node in childern:
        ans.extend(_reversed_breadth_first_flatten(graph, child_node))
    ans.append(node)
    return ans


def _print_node(graph: Graph, node, childern):
    if is_leaf(graph, node):
        return str_leaf(node)
    else:
        return str_parent(node.name, childern)


def reversed_breadth_first_flatten(graph: Graph, node: protos.axiomic.AxiomicNode) -> list:
    childern = _reversed_breadth_first_flatten(graph, node)
    deduped_childern = dedup_childern_list(childern)
    return deduped_childern


def _print_tree(graph: Graph, node: protos.axiomic.AxiomicNode):
    print_nodes = build_print_nodes(graph, node)

    for print_node in print_nodes:
        print_tree_node(print_node)


def extract_props_from_json(j: dict) -> dict:
    '''
    Convert from a json object to a list of properties. 
    Recusive pull out any string key/value pairs which are not names.
    Names are handled as childern nodes.
    '''
    props = {}
    for k, v in j.items():
        if k == 'name':
            continue
        if isinstance(v, dict):
            props.update(extract_props_from_json(v))
        if not isinstance(v, dict) and not isinstance(v, list):
            props[k] = str(v)
    return props


def get_node_props(node: protos.axiomic.AxiomicNode) -> dict:
    props = extract_props_from_json(json.loads(protos.json_dump(node)))
    # Detect empty immediate nodes and insert a special empty string place holder
    if node.WhichOneof('node_type') == 'immediate_node' and 'value' not in props:
        props['value'] = '[underline]<Empty String>[/underline]'
    return props


def __print_tree(graph: Graph, node: protos.axiomic.AxiomicNode):
    nodes = reversed_breadth_first_flatten(graph, node)

    console = Console()
    table = Table(show_header=False, header_style="dim", box=None, title=f"Graph: {node.name}", title_style="dim")

    table.add_column("Node Name", style="dim", justify="right")
    table.add_column("Node Value", style="dim")
    for node in nodes:
        childern = list_childern(graph, node)
        name, stuff = _print_node(graph, node, childern)
        table.add_row(name, stuff)
    console.print("\n")
    console.print(table)
    console.print("\n")


def node_to_print_node(graph: Graph, node: protos.axiomic.AxiomicNode):
    print_name = node.name
    props = get_node_props(node)
    childern = []

    for child_name, child_node in list_childern(graph, node):
        child_print_node = node_to_print_node(graph, child_node)
        childern.append(TreeChild(child_name, child_print_node))

    return TreePrintNode(print_name, props, childern)


def build_print_nodes(graph: Graph, node: protos.axiomic.AxiomicNode):
    nodes = reversed_breadth_first_flatten(graph, node)

    print_nodes = []
    for node in nodes:
        # childern = list_childern(graph, node)
        # name, stuff = _print_node(graph, node, childern)
        print_nodes.append(node_to_print_node(graph, node))

    return print_nodes


@dataclasses.dataclass
class TreeChild:
    relation: str
    child_node: any


@dataclasses.dataclass
class TreePrintNode:
    name: str
    props: dict
    children: list[TreeChild]

    
def node_short_desc(child_node: TreeChild):
    if len(child_node.child_node.children) == 0:
        if 'value' in child_node.child_node.props:
            panel = Panel(child_node.child_node.props["value"], 
                          title=f'{child_node.relation}: {child_node.child_node.name}',
                          title_align='left',
                          expand=False)
            return panel
            # return f'{child_node.relation}: {child_node.child_node.name} - {child_node.child_node.props["value"]}'
    return f'{child_node.relation}: {child_node.child_node.name}'


def recur_add_childern(child_node, wip_tree, max_depth=4):

    for gchild in child_node.children:
        gchild_relation = gchild.relation
        gchild_name = gchild.child_node.name
        gchild_tree = wip_tree.add(node_short_desc(gchild))
        
        if max_depth > 0:
            recur_add_childern(gchild.child_node, gchild_tree, max_depth - 1)
        else:
            gchild_tree.add('...')


def print_tree_node(node: TreePrintNode):
    console = Console()
    #console.print(f"[bold]{node.name}[/bold]")

    if len(node.props) > 0:
        props_table = Table(show_header=False, header_style="", show_lines=True, show_edge=True,
                            title=node.name, expand=False)
        props_table.add_column("Property")
        props_table.add_column("Value")
        for k, v in sorted(node.props.items()):
            props_table.add_row(str(k), str(v))
        # console.print(props_table)
        tree = Tree(props_table)
    else:
        tree = Tree(node.name)
    # tree.add(props_table)
    for child in node.children:
        ctree = tree.add(node_short_desc(child))
        recur_add_childern(child.child_node, ctree, max_depth=2)
    console.print(tree)


if __name__ == '__main__':

    node = TreePrintNode('root', {'a': 1, 'b': 2}, [TreeChild('input', 'child1'), TreeChild('output', 'child2')])
    print_tree_node(node)
    
