"""
https://github.com/jeffrichardchemistry/molbokeh/blob/main/MolBokeh/MolBokeh.py
"""

import base64
from typing import Union

import pandas as pd
from bokeh.models import ColumnDataSource, HoverTool
from rdkit import Chem
from rdkit.Chem import rdDepictor
from rdkit.Chem.Draw import rdMolDraw2D


class MolBokeh:
    def __init__(self):
        self._dfwithimg = None

    def add_molecule(
        self,
        fig,
        source: ColumnDataSource,
        smilesColName: str,
        hoverAdditionalInfo: Union[None, list] = None,
        molSize: tuple = (150, 150),
    ):
        """
        Add molecules images inside bokeh charts.

        Args:
            fig (bokeh object): Bokeh chart.
            source (ColumnDataSource): Bokehs data type used to built chart.
            smilesColName (str): Columns name that contain smiles.
            hoverAdditionalInfo (None | list): List of colnames that will show in hover.
            molSize (tuple): Size of molecules img.
        """

        UsersDF = pd.DataFrame(source.data)  # geting users df as dataframe
        dfwithMolImg = self.__makingMoleculesImg(
            df=UsersDF, smilesColName=smilesColName, molSize=molSize
        )
        fig_ = self.__cleanOriginalHover(fig=fig)
        source.data = dfwithMolImg.copy()  # update original data with base64 imgs
        hoverhtml = self.__makingHTML2hover(
            molimgColName="Mol_IMGSVG",
            hoverAdditionalInfo=hoverAdditionalInfo,
            molSize=molSize,
        )
        fig_ = self.__applyHover2fig(fig=fig_, hoverhtml=hoverhtml)

        self._dfwithimg = dfwithMolImg.copy()
        return fig_

    def __applyHover2fig(self, fig, hoverhtml: str):
        """
        Apply hover inside bokeh charts.
        """
        fig_ = fig
        hover = HoverTool(tooltips=hoverhtml)
        fig_.add_tools(hover)
        return fig_

    def __makingMoleculesImg(
        self, df: pd.DataFrame, smilesColName: str, molSize: tuple
    ):
        """
        This function creates the images and encodes them as a base64 string, it is added as a new column of the dataframe.
        """
        df_ = df.copy()
        df_["Mol_IMGSVG"] = df_[smilesColName].apply(
            lambda x: MolBokeh.smiTosvg(x, toHTMLformat=False, molSize=molSize)
        )
        return df_

    def __makingHTML2hover(
        self,
        molimgColName: str,
        hoverAdditionalInfo: Union[None, int] = None,
        molSize: tuple = (150, 150),
    ):
        """
        This function creates the HTML code that will be used as the bokeh hover.
        """
        if hoverAdditionalInfo == None:
            tooltips = f"""
            <div>
                <img src="@{molimgColName}" alt="Imagem" style="width: {molSize[0]}px; height: {molSize[0]}px;">
            </div>
            """
        else:
            tagsinfo2html = "\n".join(
                [
                    f'<p style="text-align: center;"><strong>{colnames}:</strong> @openbrac{colnames}closebrac</p>'
                    for colnames in hoverAdditionalInfo
                ]
            )
            tooltips = f"""
            <div>
                <img src="@{molimgColName}" alt="Imagem" style="width: {molSize[0]}px; height: {molSize[0]}px;">
                {tagsinfo2html}                
            </div>
            """.replace("openbrac", "{").replace("closebrac", "}")

        return tooltips

    def __cleanOriginalHover(self, fig):
        """
        This function deletes the original hover if it exists.
        """
        fig_ = fig
        for tool in fig_.tools:
            if isinstance(tool, HoverTool):
                fig_.tools.remove(tool)

        return fig_

    @staticmethod
    def smiTosvg(smi, molSize=(320, 320), kekulize=True, toHTMLformat: bool = False):
        """
        This function generate a molecule img from smiles and return as html tag or pure base64 format.
        """
        mol = Chem.MolFromSmiles(smi)
        mc = Chem.Mol(mol.ToBinary())
        if kekulize:
            try:
                Chem.Kekulize(mc)
            except:
                mc = Chem.Mol(mol.ToBinary())
        if not mc.GetNumConformers():
            rdDepictor.Compute2DCoords(mc)
        drawer = rdMolDraw2D.MolDraw2DSVG(molSize[0], molSize[1])
        drawer.DrawMolecule(mc)
        drawer.FinishDrawing()
        svg = drawer.GetDrawingText().replace("svg:", "")

        b64 = base64.b64encode(svg.encode("utf-8")).decode("utf-8")
        if toHTMLformat:
            html = r'<img src="data:image/svg+xml;base64,%s"/>' % b64
        else:
            html = f"data:image/svg+xml;base64,{b64}"

        return html
