from collections import defaultdict
from typing import Iterable, Optional
from toolz import groupby
from attr import define

from atopile import address
from atopile.address import AddrStr
from atopile.datatypes import Ref
from atopile.instance_methods import (
    all_descendants,
    get_children,
    get_links,
    get_parent,
    iter_parents,
    match_interfaces,
    match_pins_and_signals,
    match_signals,
    match_modules,
)
from atopile.loop_soup import LoopSoup
from atopile.address import get_name, add_instance


def get_nets(root: AddrStr) -> Iterable[Iterable[str]]:
    """Find all the nets under a given root."""
    net_soup = LoopSoup()
    for addr in all_descendants(root):
        if match_pins_and_signals(addr):
            net_soup.add(addr)
        for source, target in get_links(addr):
            if match_interfaces(source) and match_interfaces(target):
                for int_pin in get_children(source):
                    if match_pins_and_signals(int_pin):
                        net_soup.join(int_pin, add_instance(target, get_name(int_pin)))
                    else:
                        raise NotImplementedError
            elif match_interfaces(source) or match_interfaces(source):
                raise NotImplementedError
            else:
                net_soup.join(source, target)
    return net_soup.groups()


@define
class _Net:
    nodes_on_net: list[str]

    base_name: Optional[str] = None

    lcp: Optional[str] = None
    prefix: Optional[Ref] = None

    suffix: Optional[int] = None

    def get_name(self) -> str:
        """Get the name of the net."""
        # suffix is a ref (tuple of strings)
        return (
            f"{'-'.join(self.prefix) + '-' if self.prefix else ''}"
            f"{self.base_name or 'net'}"
            f"{'-' + str(self.suffix) if self.suffix else ''}"
        )

    def generate_base_net_name(self) -> Optional[str]:
        """TODO:"""
        WEIGHT_NO_GRANDPARENTS = 10
        WEIGHT_INTERFACE_GRANDPARENT = 5
        WEIGHT_SIGNAL = 2

        name_candidates = defaultdict(int)

        for signal in filter(match_signals, self.nodes_on_net):
            name = get_name(signal)
            if get_parent(signal) is None:
                name_candidates[name] += WEIGHT_NO_GRANDPARENTS
            elif any(map(match_interfaces, iter_parents(signal))):
                name_candidates[name] += WEIGHT_INTERFACE_GRANDPARENT
            else:
                name_candidates[name] += WEIGHT_SIGNAL

        if name_candidates:
            highest_rated_name = max(name_candidates, key=name_candidates.get)
            self.base_name = highest_rated_name


def _find_net_names(nets: Iterable[Iterable[str]]) -> dict[str, list[str]]:
    """Find the names of the nets."""
    # make net objects
    net_objs = [_Net(list(net)) for net in nets]

    # grab all the nets base names
    for net in net_objs:
        net.generate_base_net_name()

    # for the net objects that still conflict, grab a prefix
    conflicing_nets = _find_conflicts(net_objs)
    _add_prefix(conflicing_nets)

    # if they still conflict, slap a suffix on that bad boi
    conflicing_nets = _find_conflicts(net_objs)
    _add_suffix(conflicing_nets)

    return {net.get_name(): net.nodes_on_net for net in net_objs}


def _find_conflicts(nets: Iterable[_Net]) -> Iterable[Iterable[_Net]]:
    """"""
    nets_grouped_by_name = groupby(lambda net: net.get_name(), nets)
    for nets in nets_grouped_by_name.values():
        if len(nets) > 1:
            yield nets


def _add_prefix(conflicts: Iterable[list[_Net]]):
    """Resolve conflicts in net names."""
    for conflict_nets in conflicts:
        for net in conflict_nets:
            if net.base_name:
                # Find the parent of the net that is a module
                parent_module_iter = filter(
                    match_modules, iter_parents(net.nodes_on_net[0])
                )

                # Get the first parent module that matches, or None if there's no match
                parent_module = next(parent_module_iter, None)

                # Check if a parent module was found
                if parent_module:
                    # Get the ref of the parent module
                    if hasattr(parent_module, "ref"):
                        net.prefix = address.get_instance_section(parent_module)


def _add_suffix(conflicts: Iterable[list[_Net]]):
    """Add an integer suffix to the nets to resolve conflicts."""
    for conflict_nets in conflicts:
        for i, net in enumerate(conflict_nets):
            net.suffix = i


class NetFinder:
    def __init__(self) -> None:
        self.net_name_to_nodes_map: dict[AddrStr, dict[str, Iterable[AddrStr]]] = {}
        self.node_to_net_name: dict[AddrStr, dict[AddrStr, str]] = {}

    def get_nets_by_name(self, entry: AddrStr) -> dict[str, list[AddrStr]]:
        """Get the nets for a given root."""
        if address.get_instance_section(entry):
            raise ValueError("Only entry are supported for now")

        if entry not in self.net_name_to_nodes_map:
            self.net_name_to_nodes_map[entry] = _find_net_names(get_nets(entry))
            self.node_to_net_name[entry] = {
                node: net_name
                for net_name, nodes in self.net_name_to_nodes_map[entry].items()
                for node in nodes
            }

        return self.net_name_to_nodes_map[entry]

    def get_net_name_node_is_on(self, node: AddrStr) -> str:
        """Get the net name for a given node."""
        entry = address.get_entry(node)
        if entry not in self.node_to_net_name:
            self.get_nets_by_name(entry)

        return self.node_to_net_name[entry][node]


net_finder = NetFinder()


def get_net_name_node_is_on(addr: str) -> str:
    """
    Return the net name for the given address
    """
    return net_finder.get_net_name_node_is_on(addr)


def get_nets_by_name(addr: AddrStr) -> dict[str, list[AddrStr]]:
    """
    Return a dict of net names to nodes
    """
    return net_finder.get_nets_by_name(addr)
