""" Drawing utilites for each type of entry.
"""
from __future__ import annotations

from typing import Type, Optional

from PIL import Image
import rustworkx.visualization as rx_vis

import chemrecon.entrygraph.entrygraph
from chemrecon.schema import Entry, Relation

attr_none: str = "''"
linewidth: int = 30


# noinspection PyProtectedMember
class EntryGraphDrawer():
    entrygraph_type: Type[chemrecon.entrygraph.entrygraph.EntryGraph]

    # Draw settings
    draw_reconids: bool

    def __init__(
        self,
        # Settings
        draw_reconids: bool = False
    ):
        """ The initializer for the Drawer sets the settings
        """
        self.draw_reconids = draw_reconids

    def draw(
        self,
        entrygraph: chemrecon.entrygraph.entrygraph.EntryGraph,
        filename: str = None,
        filetype: str = 'png',
        scores: Optional[dict[Entry, float]] = None
    ) -> Image.Image:
        """ Draw the entry graph.

            For possible file formats, see the the corresponding
            [RustWorkX documentation](https://www.rustworkx.org/apiref/rustworkx.visualization.graphviz_draw.html).
        """

        g_ = entrygraph.g.copy()

        # Empty score dict if not given
        score_dict: dict[Entry, float] = scores if scores is not None else dict()

        # Remove duplicate edges (for symmetric relations)
        for e_index in g_.edge_indices():
            data = g_.get_edge_data_by_index(e_index)
            if data.relation.symmetric:
                x, y = g_.get_edge_endpoints_by_index(e_index)
                if x > y:
                    g_.remove_edge_from_index(e_index)

        # Remove self-edges
        for e_index in g_.edge_indices():
            x, y = g_.get_edge_endpoints_by_index(e_index)
            if x == y:
                g_.remove_edge_from_index(e_index)

        return rx_vis.graphviz_draw(
            g_,
            node_attr_fn = lambda v: self.node_attr(v.entry, score_dict.get(v.entry, None)),
            edge_attr_fn = lambda e: self.edge_attr(e.relation),
            graph_attr = None,  # TODO
            filename = filename,
            image_type = filetype,
        )

    def node_attr(self, entry: Entry, score: Optional[float]) -> dict[str, str]:
        assert entry.recon_id is not None

        # Label
        if self.draw_reconids:
            recon_id_part = f'\n{entry.recon_id}'
        else:
            recon_id_part = ''

        # Break lines of label
        entrylabel = entry._vis_str()
        label_lines = entrylabel.splitlines()
        split_label: list[str] = list()
        for l in label_lines:
            while l != '':
                split_label.append(l[:linewidth])
                l = l[linewidth:]
        entry_label_final = ('\n').join(split_label)

        # Determine label
        label: str = f'{entry.get_table_name()}\n{entry_label_final}{recon_id_part}'
        if score is not None:
            label += f'\n{score:.3f}'

        # Other drawing
        d = {
            'style': 'filled',
            'shape': 'box',
            'label': label
        }
        d.update(entry._vis_attrs())
        return d

    def edge_attr(self, rel: Relation) -> dict[str, str]:
        label = rel._vis_str()
        d = {
            'label': label,
            'dir': 'none' if rel.symmetric else 'forward',
        }

        # Update with rel_type-specific attributes and return
        d.update(rel._vis_attrs())
        return d

# Default drawer - this is used by the EntryGraph .draw() and .show() methods.
defaultdrawer = EntryGraphDrawer()
