from abc import ABC, abstractmethod
from copy import copy
import json
import logging
import os

import coreir as pycoreir

from magma.digital import Digital
from magma.array import Array
from magma.bits import Bits
from magma.backend.coreir.coreir_utils import (
    attach_debug_info, check_magma_interface, constant_to_value, get_inst_args,
    get_module_of_inst, magma_interface_to_coreir_module_type,
    magma_port_to_coreir_port, make_cparams, map_genarg,
    magma_name_to_coreir_select, Slice)
from magma.compile_exception import UnconnectedPortException
from magma.interface import InterfaceKind
from magma.is_definition import isdefinition
from magma.logging import root_logger
from magma.passes import dependencies
from magma.tuple import Tuple
from magma.backend.util import get_codegen_debug_info
from magma.clock import is_clock_or_nested_clock
from magma.wire_clock import wire_default_clock, get_default_clocks
from magma.config import get_debug_mode
from magma.protocol_type import MagmaProtocol, MagmaProtocolMeta
from magma.ref import PortViewRef, ArrayRef
from magma.symbol_table import SYMBOL_TABLE_EMPTY


# NOTE(rsetaluri): We do not need to set the level of this logger since it has
# already been done in backend/coreir/coreir_backend.py.
_logger = root_logger().getChild("coreir_backend")


_generator_callbacks = {}


def _is_generator(ckt_or_inst):
    return ckt_or_inst.coreir_genargs is not None


def _coreir_longname(magma_defn_or_decl, coreir_module_or_generator):
    # NOTE(rsetaluri): This is a proxy to exposing a pycoreir/coreir-c API to
    # get a module's longname. This logic should be identical right now. Another
    # caveat is that we don't elaborate the CoreIR generator at the magma level,
    # so it's longname needs to be dynamically reconstructed anyway.
    namespace = coreir_module_or_generator.namespace.name
    prefix = "" if namespace == "global" else f"{namespace}_"
    longname = prefix + coreir_module_or_generator.name
    if isinstance(coreir_module_or_generator, pycoreir.Module):
        return longname
    assert isinstance(coreir_module_or_generator, pycoreir.Generator)
    param_keys = coreir_module_or_generator.params.keys()
    for k in param_keys:
        v = magma_defn_or_decl.coreir_genargs[k]
        longname += f"__{k}{v}"
    return longname


def _collect_drivers(value):
    """
    Iterate over value to collect the child drivers, packing slices together
    """
    drivers = []
    start_idx = 0
    for i in range(1, len(value)):
        # If the next value item is not a reference to an array of bits where
        # the array matches the previous item and the index is incremented by
        # one, append the current slice to drivers (may introduce slices of
        # length 1)
        if not (
            isinstance(value[i].name, ArrayRef) and
            issubclass(value[i].name.array.T, Digital) and
            isinstance(value[i - 1].name, ArrayRef) and
            value[i].name.array is value[i - 1].name.array and
            value[i].name.index == value[i - 1].name.index + 1
        ):
            drivers.append(value[start_idx:i])
            start_idx = i
    drivers.append(value[start_idx:])
    return drivers


def _unwrap(x):
    if isinstance(x, MagmaProtocol):
        return x._get_magma_value_()
    if isinstance(x, MagmaProtocolMeta):
        return x._to_magma_()
    return x


class TransformerBase(ABC):
    __MISSING = object()

    def __init__(self, backend, opts):
        self.backend = backend
        self.opts = opts
        self.ran = False
        self._children = None

    def run(self):
        if self.ran:
            raise RuntimeError("Can only run transformer once")
        self._children = self.children()
        for child in self._children:
            child.run()
        self.run_self()
        self.ran = True

    @abstractmethod
    def children(self):
        raise NotImplementedError()

    def run_self(self):
        pass

    def get_opt(self, key, default=__MISSING):
        if default is TransformerBase.__MISSING:
            return self.opts[key]
        return self.opts.get(key, default)


class LeafTransformer(TransformerBase):
    def children(self):
        return []


class DefnOrDeclTransformer(TransformerBase):
    def __init__(self, backend, opts, defn_or_decl):
        super().__init__(backend, opts)
        self.defn_or_decl = defn_or_decl
        self.coreir_module = None

    def children(self):
        if _is_generator(self.defn_or_decl):
            return [GeneratorTransformer(
                self.backend, self.opts, self.defn_or_decl)]
        try:
            coreir_module = self.backend.get_module(self.defn_or_decl)
            _logger.debug(f"{self.defn_or_decl} already compiled, skipping")
            self.coreir_module = coreir_module
            return []
        except KeyError:
            pass
        if not isdefinition(self.defn_or_decl):
            return [DeclarationTransformer(self.backend,
                                           self.opts,
                                           self.defn_or_decl)]
        wrapped = getattr(self.defn_or_decl, "wrappedModule", None)
        if wrapped and wrapped.context is self.backend.context:
            return [WrappedTransformer(self.backend,
                                       self.opts,
                                       self.defn_or_decl)]
        return [DefinitionTransformer(self.backend,
                                      self.opts,
                                      self.defn_or_decl)]

    def run_self(self):
        self._run_self_impl()
        self._generate_symbols()

    def _generate_symbols(self):
        if not self.get_opt("generate_symbols", False):
            return
        out_module_name = _coreir_longname(
            self.defn_or_decl, self.coreir_module)
        self.opts.get("symbol_table").set_module_name(
            self.defn_or_decl.name, out_module_name)

    def _run_self_impl(self):
        if self.coreir_module:
            return
        self.coreir_module = self._children[0].coreir_module
        self.backend.add_module(self.defn_or_decl, self.coreir_module)
        if isdefinition(self.defn_or_decl):
            self.defn_or_decl.wrappedModule = self.coreir_module
            libs = self.backend.included_libs()
            self.defn_or_decl.coreir_wrapped_modules_libs_used = libs


class GeneratorTransformer(TransformerBase):
    def __init__(self, backend, opts, defn_or_decl):
        super().__init__(backend, opts)
        self.defn_or_decl = defn_or_decl
        self.coreir_module = None

    def children(self):
        try:
            coreir_module = self.backend.get_module(self.defn_or_decl)
            _logger.debug(f"{self.defn_or_decl} already compiled, skipping")
            self.coreir_module = coreir_module
            return []
        except KeyError:
            pass
        assert not isdefinition(self.defn_or_decl)
        return [DeclarationTransformer(self.backend,
                                       self.opts,
                                       self.defn_or_decl)]

    def run_self(self):
        self._generate_symbols()
        if self.coreir_module is not None:
            return
        self.coreir_module = self._children[0].coreir_module

    def _generate_symbols(self):
        if not self.get_opt("generate_symbols", False):
            return
        global _generator_callbacks

        def _callback(coreir_inst):
            magma_names = list(self.defn_or_decl.interface.ports.keys())
            coreir_names = list(k for k, _ in coreir_inst.module.type.items())
            assert len(magma_names) == len(coreir_names)
            for magma_name, coreir_name in zip(magma_names, coreir_names):
                self.opts.get("symbol_table").set_port_name(
                    self.defn_or_decl.name, magma_name, coreir_name)

        assert self.defn_or_decl not in _generator_callbacks
        _generator_callbacks[self.defn_or_decl] = _callback


class InstanceTransformer(LeafTransformer):
    def __init__(self, backend, opts, inst, defn):
        super().__init__(backend, opts)
        self.inst = inst
        self.defn = defn
        self.coreir_inst_gen = None

    def run_self(self):
        self.coreir_inst_gen = self.run_self_impl()

    def run_self_impl(self):
        _logger.debug(
            f"Compiling instance {(self.inst.name, type(self.inst).name)}"
        )
        defn = type(self.inst)
        if hasattr(self.inst, "namespace"):
            lib = self.backend.get_lib(self.inst.namespace)
        else:
            lib = self.backend.get_lib(self.inst.coreir_lib)
            if self.inst.coreir_lib == "global":
                lib = self.get_opt("user_namespace", lib)
        if not _is_generator(self.inst):
            module = get_module_of_inst(self.backend.context, self.inst, lib)
            args = get_inst_args(self.inst)
            args = self.backend.context.new_values(args)
            return lambda m: m.add_module_instance(self.inst.name, module, args)
        generator = lib.generators[defn.coreir_name]
        config_args = {k: v for k, v in self.inst.coreir_configargs.items()}
        config_args = self.backend.context.new_values(config_args)
        gen_args = {k: map_genarg(self.backend.context, v)
                    for k, v in defn.coreir_genargs.items()}
        gen_args = self.backend.context.new_values(gen_args)
        return lambda m: m.add_generator_instance(
            self.inst.name, generator, gen_args, config_args)


class WrappedTransformer(LeafTransformer):
    def __init__(self, backend, opts, defn):
        super().__init__(backend, opts)
        self.defn = defn
        self.coreir_module = self.defn.wrappedModule
        self.backend.include_lib_or_libs(
            self.defn.coreir_wrapped_modules_libs_used)


class DefinitionTransformer(TransformerBase):
    def __init__(self, backend, opts, defn):
        super().__init__(backend, opts)
        self.defn = defn
        self.coreir_module = None
        self.decl_tx = DeclarationTransformer(self.backend,
                                              self.opts,
                                              self.defn)
        self.inst_txs = {
            inst: InstanceTransformer(self.backend, self.opts, inst, self.defn)
            for inst in self.defn.instances
        }
        self.clocks = get_default_clocks(defn)
        self._constant_cache = {}

    def children(self):
        children = []
        if not self.get_opt("skip_instance_graph", False):
            deps = dependencies(self.defn, include_self=False)
            opts = self.opts.copy()
            opts.update({"skip_instance_graph": True})
            children += [DefnOrDeclTransformer(self.backend, opts, dep)
                         for dep in deps]
        children += [self.decl_tx]
        children += self.inst_txs.values()
        return children

    def run_self(self):
        _logger.debug(f"Compiling definition {self.defn}")
        self.coreir_module = self.decl_tx.coreir_module
        if self.defn.inline_verilog_strs:
            inline_verilog = "\n\n".join(x[0] for x in
                                         self.defn.inline_verilog_strs)
            connect_references = {}
            for _, inline_value_map in self.defn.inline_verilog_strs:
                for key, value in inline_value_map.items():
                    connect_references[key] = magma_port_to_coreir_port(value)
            self.coreir_module.add_metadata("inline_verilog", json.dumps(
                {"str": inline_verilog,
                 "connect_references": connect_references}
            ))
        for name, module in self.defn.compiled_bind_modules.items():
            self.backend.bind_module(name, module)

        self.coreir_module.definition = self.get_coreir_defn()

    def _generate_symbols(self, coreir_insts):
        if not self.get_opt("generate_symbols", False):
            return
        for inst, coreir_inst in coreir_insts.items():
            self.get_opt("symbol_table").set_instance_name(
                self.defn.name, inst.name,
                (SYMBOL_TABLE_EMPTY, coreir_inst.name))
            self.get_opt("symbol_table").set_instance_type(
                self.defn.name, inst.name, type(inst).name)

    def get_coreir_defn(self):
        coreir_defn = self.coreir_module.new_definition()
        coreir_insts = {inst: self.inst_txs[inst].coreir_inst_gen(coreir_defn)
                        for inst in self.defn.instances}
        # Call generator callback if necessary.
        global _generator_callbacks
        for inst, coreir_inst in coreir_insts.items():
            try:
                callback = _generator_callbacks.pop(type(inst))
            except KeyError:
                continue
            callback(coreir_inst)
        self._generate_symbols(coreir_insts)
        # If this module was imported from verilog, do not go through the
        # general module construction flow. Instead just attach the verilog
        # source as metadata and return the module.
        if hasattr(self.defn, "verilogFile") and self.defn.verilogFile:
            metadata = json.dumps({"verilog_string": self.defn.verilogFile})
            self.coreir_module.add_metadata("verilog", metadata)
            return coreir_defn
        if hasattr(self.defn, "verilog") and self.defn.verilog:
            metadata = json.dumps({"verilog_body": self.defn.verilog})
            self.coreir_module.add_metadata("verilog", metadata)
            return coreir_defn
        if self.defn.coreir_lib is not None:
            self.backend.include_lib_or_libs(self.defn.coreir_lib)
        for name, port in self.defn.interface.ports.items():
            _logger.debug(f"{name}, {port}, {port.is_output()}")
        for inst, coreir_inst in coreir_insts.items():
            if get_codegen_debug_info() and getattr(inst, "debug_info", False):
                attach_debug_info(coreir_inst, inst.debug_info)
            if getattr(inst, "coreir_metadata"):
                for k, v in inst.coreir_metadata.items():
                    coreir_inst.add_metadata(k, json.dumps(v))
        for inst in coreir_insts:
            for name, port in inst.interface.ports.items():
                self.connect_non_outputs(coreir_defn, port)
        for port in self.defn.interface.ports.values():
            self.connect_non_outputs(coreir_defn, port)
        return coreir_defn

    def connect_non_outputs(self, module_defn, port):
        # Recurse into non input types that may contain inout children.
        if isinstance(port, Tuple) and not port.is_input() or \
           isinstance(port, Array) and not port.T.is_input():
            for elem in port:
                self.connect_non_outputs(module_defn, elem)
        elif not port.is_output():
            self.connect(module_defn, port, port.trace())

    def get_source(self, port, value, module_defn):
        port = _unwrap(port)
        value = _unwrap(value)
        if isinstance(value, pycoreir.Wireable):
            return value
        if isinstance(value, Slice):
            return module_defn.select(value.get_coreir_select())
        if isinstance(value, Bits) and value.const():
            return self._const_instance(value, len(value), module_defn)
        if value.anon() and isinstance(value, Array):
            drivers = _collect_drivers(value)
            offset = 0
            for d in drivers:
                d = _unwrap(d)
                if len(d) == 1:
                    # _collect_drivers will introduce a slice of length 1 for
                    # non-slices, so we index them here with 0 to unpack the
                    # extra array dimension
                    self.connect(module_defn, port[offset], d[0])
                else:
                    self.connect(module_defn,
                                 Slice(port, offset, offset + len(d)),
                                 Slice(d[0].name.array, d[0].name.index,
                                       d[-1].name.index + 1))
                offset += len(d)

            return None
        if isinstance(value, Tuple) and value.anon():
            for p, v in zip(port, value):
                self.connect(module_defn, p, v)
            return None
        if value.const():
            return self._const_instance(value, None, module_defn)
        if isinstance(value.name, PortViewRef):
            return module_defn.select(
                magma_name_to_coreir_select(value.name))
        return module_defn.select(magma_port_to_coreir_port(value))

    def connect(self, module_defn, port, value):
        if value is None and is_clock_or_nested_clock(type(port)):
            if not wire_default_clock(port, self.clocks):
                # No default clock
                raise UnconnectedPortException(port)
            value = port.trace()
        if value is None:
            if port.is_inout():
                return  # skip inouts because they might be conn. as an input.
            if getattr(self.defn, "_ignore_undriven_", False):
                return
            raise UnconnectedPortException(port)
        source = self.get_source(port, value, module_defn)
        if not source:
            return
        sink = module_defn.select(magma_port_to_coreir_port(port))
        module_defn.connect(source, sink)
        if get_codegen_debug_info() and getattr(port, "debug_info", False):
            attach_debug_info(module_defn, port.debug_info, source, sink)

    def _const_instance(self, constant, num_bits, module_defn):
        value = constant_to_value(constant)
        key = (value, num_bits)
        try:
            return self._constant_cache[key]
        except KeyError:
            pass
        if num_bits is None:
            config = self.backend.context.new_values({"value": bool(value)})
            name = f"bit_const_{value}_{num_bits}"
            mod = self.backend.get_lib("corebit").modules["const"]
            module_defn.add_module_instance(name, mod, config)
        else:
            config = self.backend.context.new_values({"value": value})
            name = f"const_{value}_{num_bits}"
            gen = self.backend.get_lib("coreir").generators["const"]
            gen_args = self.backend.context.new_values({"width": num_bits})
            module_defn.add_generator_instance(name, gen, gen_args, config)
        out = module_defn.select(f"{name}.out")
        return self._constant_cache.setdefault(key, out)


class DeclarationTransformer(LeafTransformer):
    def __init__(self, backend, opts, decl):
        super().__init__(backend, opts)
        self.decl = decl
        self.coreir_module = None

    def run_self(self):
        self.coreir_module = self._run_self_impl()
        self._generate_symbols()

    def _generate_symbols(self):
        if not self.get_opt("generate_symbols", False):
            return
        if _is_generator(self.decl):
            return
        magma_names = list(self.decl.interface.ports.keys())
        coreir_names = list(k for k, _ in self.coreir_module.type.items())
        assert len(magma_names) == len(coreir_names)
        for magma_name, coreir_name in zip(magma_names, coreir_names):
            self.opts.get("symbol_table").set_port_name(
                self.decl.name, magma_name, coreir_name)

    def _run_self_impl(self):
        self.decl = self.decl
        _logger.debug(f"Compiling declaration {self.decl}")
        if self.decl.coreir_lib is not None:
            self.backend.include_lib_or_libs(self.decl.coreir_lib)
        # These libraries are already available by default in coreir, so we
        # don't need declarations.
        if self.decl.coreir_lib in ["coreir", "corebit", "commonlib",
                                    "memory"]:
            lib = self.backend.get_lib(self.decl.coreir_lib)
            if not _is_generator(self.decl):
                return lib.modules[self.decl.coreir_name]
            return lib.generators[self.decl.coreir_name]
        try:
            coreir_module = self.backend.get_module(self.decl)
            _logger.debug(f"{self.decl} already compiled, skipping")
            return coreir_module
        except KeyError:
            pass
        if get_debug_mode():
            check_magma_interface(self.decl.interface)
        module_type = magma_interface_to_coreir_module_type(
            self.backend.context, self.decl.interface)
        if isinstance(self.decl.interface, InterfaceKind):
            module_type = self.backend.context.Flip(module_type)
        kwargs = {}
        if hasattr(self.decl, "coreir_config_param_types"):
            param_types = self.decl.coreir_config_param_types
            kwargs["cparams"] = make_cparams(self.backend.context, param_types)
        if hasattr(self.decl, "namespace"):
            # Allows users to choose namespace explicitly with
            # class MyCircuit(m.Circuit):
            #     namespace = "foo"
            # overrides user_namespace setting
            namespace = self.backend.get_lib(self.decl.namespace)
        else:
            namespace = self.get_opt("user_namespace",
                                     self.backend.context.global_namespace)
        coreir_module = namespace.new_module(
            self.decl.coreir_name, module_type, **kwargs)
        if get_codegen_debug_info() and self.decl.debug_info:
            attach_debug_info(coreir_module, self.decl.debug_info)
        for key, value in self.decl.coreir_metadata.items():
            coreir_module.add_metadata(key, json.dumps(value))
        return coreir_module
