#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Copyright 2020-2022 F4PGA Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0
"""
This is intended to provide a range of helper functions around the output of
Yosys' `write_json`. Depending on the tasks, this may need to be flattened
and/or in AIG format. In any case, at minimum `proc` and usually `prep` should
be used before outputting the JSON.
"""

import json
import pprint
import re


class YosysModule:
    def __init__(self, name, module_data):
        self.name = name
        self.data = module_data

    def __str__(self):
        return "YosysModule({},\n{})".format(
            self.name, pprint.pformat(self.data))

    @property
    def ports(self):
        """List of ports on a module.

        Returns a list of tuples:
        -------
        name : str
        width : int
            The width in bits
        dir : str
            The direction, should be either `input` or `output`
        """
        plist = []
        for port, pdata in sorted(self.data["ports"].items()):
            plist.append(
                (port, len(pdata["bits"]), pdata["bits"], pdata["direction"]))
        return plist

    @property
    def cells(self):
        """List of cells of a module, excluding Yosys-internal cells
        beginning with $.

        Returns a list of tuples:
        -------
        name : str
        type : str
        """
        clist = []
        for cell, cdata in sorted(self.data["cells"].items()):
            if cdata["type"].startswith('$'):
                continue
            clist.append((cell, cdata["type"]))
        return clist

    @property
    def all_cells(self):
        """List of cells of a module, including Yosis-internal cells

        Returns a list of tuples:
        -------
        name : str
        type : str
        """
        clist = []
        for cell, cdata in sorted(self.data["cells"].items()):
            clist.append((cell, cdata["type"]))
        return clist

    @property
    def net_names(self):
        """List the net names avilable in the design."""
        return self.data["netnames"].keys()

    @property
    def nets(self):
        """List the net ids available in the design."""
        return list(
            sorted(set(n['bits'][0] for n in self.data["netnames"].values())))

    def cell_type(self, cell):
        """Return the type of a given cell"""
        for cname, cdata in self.data["cells"].items():
            if cname == cell:
                return cdata["type"]
        return None

    @property
    def module_attrs(self):
        """All attributes of a module as a dictionary"""
        return self.data["attributes"]

    def attr(self, attr, defval=None):
        """Get an attribute, or defval is not set"""
        if attr in self.module_attrs:
            return self.module_attrs[attr]
        else:
            return defval

    def __getattr__(self, attr):
        return self.attr(attr)

    def has_attr(self, attr):
        """Return true if an attribute exists"""
        return attr in self.module_attrs

    def cell_attrs(self, cell):
        """All attributes of a given cell as a dictionary"""
        return self.data["cells"][cell]["attributes"]

    def cell_attr(self, cell, attr, defval=None):
        """Get an attribute of a given cell, or defval is not set"""
        if attr in self.cell_attrs(cell):
            return self.cell_attrs(cell)[attr]
        else:
            return defval

    def net_attrs(self, netname):
        """Get all attributes of a given net as a dictionary"""
        return self.data["netnames"][netname]["attributes"]

    def net_attr(self, netname, attr, defval=None):
        """Get an attribute of a given net (specified by name),
        or defval is not set"""
        if attr in self.net_attrs(netname):
            return self.net_attrs(netname)[attr]
        else:
            return defval

    PORT_REGEX = re.compile('(.*)(\\[([0-9]+)\\])')

    def port_attrs(self, port_pin):
        """Gets all the attributes for for given pin

        Returns attributes dictionary:
        -------

        port_pin: str
        """
        m = self.PORT_REGEX.match(port_pin)
        if not m:
            return self.net_attrs(port_pin)
        port_name, _, port_idx = m.groups()
        return self.net_attrs(port_name)

    # TODO: the below code is kind of ugly, but because module and cell IO
    # specifications are inconsistent in how they are represented in the JSON,
    # it's hard to make any nicer...

    def port_conns(self, port):
        """The connections of a port

        Returns a list of connections:
        -------
        net : int
        """
        for pname, pdata in sorted(self.data["ports"].items()):
            if pname == port:
                return pdata["bits"]

    def cell_conns(self, cell, direction="input"):
        """The connections of a cell in a given direction as a 2-tuple

        Returns a list of tuples:
        -------
        port : str
        net : int
        """
        cdata = self.data["cells"][cell]
        conns = []
        for port, condata in sorted(cdata["connections"].items()):
            if cdata["port_directions"][port] == direction:
                N = len(condata)
                if N == 1:
                    conns.append((port, condata[0]))
                else:
                    for i in range(N):
                        conns.append(("{}[{}]".format(port, i), condata[i]))
        return conns

    def cell_conn_list(self, cell, port):
        """The connections of a cell in a given port as a list

        Returns a list of connections:
        -------
        net : list

        """
        data = self.data["cells"][cell]

        conns = []

        for cport, condata in sorted(data["connections"].items()):
            if cport == port:
                conns = condata

        return conns

    def cell_clk_conn(self, cell):
        """The clock net related to a given cell

        If exists returns the clock net

        net : list
        """
        data = self.data["cells"][cell]

        if "CLK" in data["connections"].keys():
            return data["connections"]["CLK"]
        else:
            return None

    def conn_io(self, net, iodir):
        """Returns a list of top level IO matching a direction and connected net number

        Returns a list:
        -------
        port : str
        """
        conn_io = []
        for port, pdata in sorted(self.data["ports"].items()):
            if pdata["direction"] == iodir:
                if net in pdata["bits"]:
                    if len(pdata["bits"]) == 1:
                        conn_io.append(port)
                    else:
                        conn_io.append(
                            "{}[{}]".format(port, pdata["bits"].index(net)))
        return conn_io

    def conn_ports(self, net, pdir):
        """Returns any cell ports matching a direction and connected net number

        Returns a list of tuples:
        -------
        cell : str
        port : str
        """
        conn_ports = []
        for cell in sorted(self.data["cells"].keys()):
            cdata = self.data["cells"][cell]
            for port, condata in sorted(cdata["connections"].items()):
                if cdata["port_directions"][port] == pdir:
                    if net in condata:
                        if len(condata) == 1:
                            conn_ports.append((cell, port))
                        else:
                            conn_ports.append(
                                (
                                    cell, "{}[{}]".format(
                                        port, condata.index(net))))
        return conn_ports

    def net_name(self, netid):
        """Get a net name from a given netid. "netid" is a design net unique id
        number.

        Returns string:
        -------
        netid : int
        """
        import sys
        print("Netid", netid, file=sys.stderr)
        names = []
        for n, props in self.data["netnames"].items():
            if netid in props['bits']:
                names.append(n)
        # FIXME: This fails when there are two wires with different names
        # connected to the same netid (yes, that's possible). Which name
        # should be returned in that case ?
        if len(names) != 1:
            raise KeyError("Net id {} not found".format(netid))
        return names[0]

    def net_attrs_by_netid(self, netid):
        """
        Returns attributes of a given netid. Raises RuntimeError if the same
        attribute is defined for two or more wires belonging to the same net.
        The attribute 'src' is an exception, 'src' strings are concatenated
        when defined for more than one wire.

        Returns dict:
        -------
        netid : int
        """

        attributes = {}

        for name, data in self.data["netnames"].items():
            if netid in data['bits']:

                # Join attributes
                for attr, value in data['attributes'].items():

                    # Allow multiple definitions of the same attribute only
                    # for the 'src'. Otherwise raise an exception
                    if attr in attributes and attr != 'src':
                        raise RuntimeError(
                            "Conflicting attributes '{}' on netid {}".format(
                                attr, netid))

                    # Join 'src' attribute strings
                    if attr in attributes and attr == 'src':
                        attributes[attr] += ";{}".format(value)
                    # Store the value
                    else:
                        attributes[attr] = value

        return attributes

    def net_drivers(self, net):
        """Returns a list of drivers of a given net, both top level inputs.
        and cell outputs. "cell" is set to the name of the module for top level
        IO.

        Returns a list of tuples:
        -------
        cell : str
        port : str
        """
        # top level *inputs* and cell outputs are both drivers
        io_drivers = [(self.name, _) for _ in self.conn_io(net, "input")]
        cell_drivers = self.conn_ports(net, "output")
        return io_drivers + cell_drivers

    def net_sinks(self, net):
        """Returns a list of sinks of a given net, both top level outputs.
        and cell inputs. "cell" is set to the name of the module for top level
        IO.

        Returns a list of tuples:
        -------
        cell : str
        port : str
        """
        # top level *outputs* and cell inputs are both drivers
        io_drivers = [(self.name, _) for _ in self.conn_io(net, "output")]
        cell_drivers = self.conn_ports(net, "input")
        return io_drivers + cell_drivers


class YosysJSON:
    def __init__(self, j, top=None):
        """Takes either the filename to a JSON file, or already parsed JSON
        data as a dictionary. Optionally the top level module can be specified
        too."""
        if (isinstance(j, str)):
            with open(j, 'r') as f:
                self.data = json.load(f)
        else:
            self.data = j
        if top is None:
            if len(self.data["modules"]) == 1:
                self.top = list(self.data["modules"].keys())[0]
            else:
                self.top = None
        else:
            self.top = top

    def module(self, module):
        """Get a given module (by name) as a `YosysModule`"""
        if module not in self.data["modules"]:
            raise KeyError(
                "No yosys module named {} (only have {})".format(
                    module, self.data["modules"].keys()))
        return YosysModule(module, self.data["modules"][module])

    def modules_with_attr(self, attr_name, attr_value):
        """Return a list of `YosysModule`s, selecting based on
        a given attribute"""
        mods = []
        for mod in self.data["modules"]:
            ymod = self.module(mod)
            if ymod.has_attr(attr_name) and ymod.attr(attr_name) == attr_value:
                mods.append(ymod)
        return mods

    def all_modules(self):
        """Return a list of the names of all modules in the design,
        sorted alphabetically"""
        return sorted(self.data["modules"].keys())

    @property
    def top_module(self):
        """Get a given module (by name) as a `YosysModule`"""
        return self.module(self.top)

    def has_module(self, module):
        """Return true if a module exists"""
        return module in self.data["modules"]

    def get_module_file(self, module):
        """Return the filename in which a given module lives"""
        src = self.module(module).attr("src")
        cpos = src.rfind(":")
        return src[0:cpos]