from lynxkite_core.ops import op
from lynxkite_graph_analytics.core import Bundle, TableName, ColumnNameByTableName
import base64
import io


def pil_to_data(image):
    buffer = io.BytesIO()
    image.save(buffer, format="webp")
    b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
    return "data:image/webp;base64," + b64


def smiles_to_data(smiles):
    import rdkit

    m = rdkit.Chem.MolFromSmiles(smiles)
    if m is None:
        return None
    img = rdkit.Chem.Draw.MolToImage(m)
    data = pil_to_data(img)
    return data


@op("LynxKite Graph Analytics", "Draw molecules", icon="microscope-filled")
def draw_molecules(
    b: Bundle,
    *,
    table_name: TableName,
    smiles_column: ColumnNameByTableName,
    image_column: str = "image",
):
    """Adds molecule images in a table."""
    b = b.copy()
    df = b.dfs[table_name]
    df = df.copy()
    df[image_column] = df[smiles_column].apply(smiles_to_data)
    b.dfs[table_name] = df
    return b