"""Add a way to visualize using Mol*."""

# =======================
# Python internal package
# =======================
# D
import dataclasses

# =======================
# Python external package
# =======================
# M
import molviewspec

# ==============
# Module package
# ==============
# U
from ..utils.generic import (
    DataClass,
    InteractionType,
    KeyIdentifier,
    PharmacophoreFamily,
)


class Molstar:
    """Generate a Mol* (MolStar) visualization from pharmacophore and
    interaction data.
    """

    def __init__(
        self,
        path: str,
        pharmacophore_collection_a: dict[PharmacophoreFamily, list[DataClass]],
        pharmacophore_collection_b: dict[PharmacophoreFamily, list[DataClass]],
        interaction_collection: dict[InteractionType, list[DataClass]],
        drawning_configuration: DataClass,
    ) -> None:
        """Build and export a Mol* visualization state.

        Parameters
        ----------
        path : `str`
            Output file path where the Mol* JSON state will be written.

        pharmacophore_collection_a : `dict[PharmacophoreFamily, list[DataClass]]`
            First set of pharmacophore features, grouped by family.

        pharmacophore_collection_b : `dict[PharmacophoreFamily, list[DataClass]]`
            Second set of pharmacophore features, grouped by family.

        interaction_collection : `dict[InteractionType, list[DataClass]]`
            Molecular interactions grouped by interaction type.

        drawning_configuration : `DataClass`
            Configuration dataclass defining visualization parameters.
        """
        self.__drawning_configuration: DataClass = drawning_configuration
        molstar_builder: molviewspec.builder.Root = (
            molviewspec.create_builder()
        )

        key_list: set = set(
            list(pharmacophore_collection_a.keys())
            + list(pharmacophore_collection_b.keys())
        )

        for key in key_list:
            pharmacophore: list = []

            if key in pharmacophore_collection_a:
                pharmacophore += pharmacophore_collection_a[key]

            if key in pharmacophore_collection_b:
                pharmacophore += pharmacophore_collection_b[key]

            primitive: molviewspec.builder.Primitives = (
                molstar_builder.primitives(opacity=0.5)
            )

            self.__draw_pharmacophore(primitive, pharmacophore, key)

        for key, interaction in interaction_collection.items():
            primitive: molviewspec.builder.Primitives = (
                molstar_builder.primitives(opacity=1)
            )

            self.__drawn_interaction(primitive, interaction, key)

        with open(path, "w", encoding="utf-8") as file:
            file.write(molstar_builder.get_state().dumps(indent=0))

    def __fetch_fields(self, key: str) -> dict:
        """Extract visualization parameters associated with a given key.

        Parameters
        ----------
        key : `str`
            Identifier prefix used to match configuration fields.

        Returns
        -------
        `dict`
            Dictionary of visualization properties extracted from the
            drawing configuration.
        """
        fetched_field: dict = {}

        for field_i in dataclasses.fields(self.__drawning_configuration):
            if not field_i.name.startswith(f"{key}__"):
                continue

            property_name: str = field_i.name.split("__")[-1]
            fetched_field[property_name] = getattr(
                self.__drawning_configuration, field_i.name
            )

        return fetched_field

    def __draw_pharmacophore(
        self,
        primitive: molviewspec.builder.Primitives,
        collection: list[DataClass],
        key: PharmacophoreFamily,
    ) -> None:
        """Draw pharmacophore features as spheres in the Mol* scene.

        Parameters
        ----------
        primitive : `molviewspec.builder.Primitives`
            Mol* primitive builder used to add graphical objects.

        collection : `list[DataClass]`
            List of pharmacophore dataclass instances to draw.

        key : `PharmacophoreFamily`
            Pharmacophore family associated with the collection.
        """
        for pharmacophore in collection:
            if pharmacophore.key != KeyIdentifier.PHARMACOPHORE:
                continue

            primitive.sphere(
                center=(
                    pharmacophore.coordinate_x,
                    pharmacophore.coordinate_y,
                    pharmacophore.coordinate_z,
                ),
                tooltip=str(key),
                **self.__fetch_fields(key),
            )

    def __drawn_interaction(
        self,
        primitive: molviewspec.builder.Primitives,
        collection: list[DataClass],
        key: InteractionType,
    ) -> None:
        """Draw molecular interactions as tubes in the Mol* scene.

        Parameters
        ----------
        primitive : `molviewspec.builder.Primitives`
            Mol* primitive builder used to add graphical objects.

        collection : `list[DataClass]`
            List of interaction dataclass instances to draw.

        key : `InteractionType`
            Interaction type associated with the collection.
        """
        for interaction in collection:
            primitive.tube(
                start=(
                    interaction.molecule_a.coordinate_x,
                    interaction.molecule_a.coordinate_y,
                    interaction.molecule_a.coordinate_z,
                ),
                end=(
                    interaction.molecule_b.coordinate_x,
                    interaction.molecule_b.coordinate_y,
                    interaction.molecule_b.coordinate_z,
                ),
                tooltip=str(key),
                **self.__fetch_fields(key),
            )
