#!/bin/env python
import newick
import json
import argparse

def read_tree( fname ):
    tree = None
    with open(fname,'r') as fh:
        nwstr = ''
        for line in fh:
            line = line.strip().rstrip('\n')
            if len(line)==0:
                continue
            if line[0]=='#':
                continue
            nwstr = nwstr+line
        try:
            tree = newick.loads(nwstr)
        except:
            tree = None
    return tree

def read_json( fname ):
    data = None
    with open(fname,'r') as fh:
        data = json.load( fh )
        # return possible wrapper
        if "data" in data:
            data = data["data"]
        return data
    return None

#------------------------------------------------------------------------------------
def node_to_dict(node):
    """Convert a newick.Node object to a nested dictionary for JSON serialization."""
    node_dict = {
        "key": node.name or "",
    }
    #if node.length is not None:
    #    node_dict["length"] = node.length
    if node.descendants:
        sorted_children = sorted(
            node.descendants,
            key=lambda x: x.name or "~"  # Tilde ensures unnamed nodes sort last
        )
        node_dict["children"] = [node_to_dict(child) for child in sorted_children]
    return node_dict

#------------------------------------------------------------------------------------
def dict_to_newick(node):
    """Recursively convert a nested dictionary to a Newick string."""
    if "children" in node and node["children"]:
        children_str = ",".join(dict_to_newick(child) for child in node["children"])
        subtree = f"({children_str})"
    else:
        subtree = ""

    name = node.get("key", "")
    length = f":{node['length']}" if "length" in node else ""

    return f"{subtree}{name}{length}"

def json_to_newick(json_obj):
    """Convert a JSON object (as dict) to a Newick string."""
    
    if isinstance(json_obj, list):
        return "\n".join(dict_to_newick(tree) + ";" for tree in json_obj)
    else:
        return dict_to_newick(json_obj) + ";"

#------------------------------------------------------------------------------------
def newick_or_json(fname):
    """Convert a Newick string into a JSON string."""
    trees = read_tree( fname )
    if trees is not None:
        # If multiple trees, return a list of dicts; otherwise, a single dict
        return json.dumps([node_to_dict(tree) for tree in trees], indent=2)
    data = read_json( fname )
    if data is not None:
        return json_to_newick( data )
    
    return None


# Example usage
if __name__ == "__main__":
    
    parser = argparse.ArgumentParser(description='Convert between newick and JSON')

    parser.add_argument('-f', '--file', dest='filename',
                        help='newick or JSON file', metavar='FILE', required=True )
    options = parser.parse_args()
    
    output = newick_or_json(options.filename)
    if output is not None:
        print(output)
