import gradio as gr
from rdkit import Chem
from rdkit.Chem import Draw
from rdkit.Chem.Draw import rdMolDraw2D

list_smiles = [
    "C[C@H](N)C(=O)O",  # - Alanine
    "CC(=O)OC1=CC=CC=C1C(=O)O",  # - Aspirin
    "CCN(CC)CC",  # - Triethylamine
    "c1ccccc1O",  # - Phenol
    "CC(C)CC(=O)O",  # - Valeric acid
    "CN1C=NC2=C1N=CN2",  # - Adenine
    "O=C(O)C1=CC=CC=C1",  # - Benzoic acid
    "C1CCCCC1",  # - Cyclohexane
    "CC(=O)N1CCCCC1",  # - (Unnamed)
]


def _extract_smiles(smiles_input):
    """
    Normalize the Gradio input into a list of non-empty SMILES strings.
    """
    if smiles_input is None:
        return []

    # Handle text areas or other string-based inputs
    if isinstance(smiles_input, str):
        return [line.strip() for line in smiles_input.splitlines() if line.strip()]

    smiles_strings = []
    for row in smiles_input:
        if isinstance(row, str):
            candidate = row.strip()
            if candidate:
                smiles_strings.append(candidate)
            continue

        if not row:
            continue
        first_cell = row[0]
        if isinstance(first_cell, str):
            candidate = first_cell.strip()
            if candidate:
                smiles_strings.append(candidate)

    return smiles_strings


def render_images(smiles_input):
    """
    Converts SMILES strings into a list of PIL Images for gallery display.
    Invalid SMILES entries are skipped.
    """
    smiles_strings = _extract_smiles(smiles_input)
    gallery_items = []

    for smiles in smiles_strings:
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            continue
        image = Draw.MolToImage(mol, size=(300, 300))
        gallery_items.append((image, smiles))

    return gallery_items


def render_svgs(smiles_input):
    """
    Converts SMILES strings into SVG representations of the molecules.
    Returns HTML containing all SVGs and error notices for invalid inputs.
    """
    smiles_strings = _extract_smiles(smiles_input)
    svg_fragments = []
    invalid_smiles = []

    for smiles in smiles_strings:
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            invalid_smiles.append(smiles)
            continue

        drawer = rdMolDraw2D.MolDraw2DSVG(300, 300)
        drawer.DrawMolecule(mol)
        drawer.FinishDrawing()
        svg = drawer.GetDrawingText().replace("svg:", "")
        svg_fragments.append(
            f"<figure style='display:inline-block;margin:0 16px 16px 0;'>"
            f"<figcaption style='text-align:center;font-family:monospace;margin-bottom:8px;'>{smiles}</figcaption>"
            f"{svg}"
            "</figure>"
        )

    html_parts = []
    if svg_fragments:
        html_parts.append("<div>" + "".join(svg_fragments) + "</div>")

    if invalid_smiles:
        invalid_list_items = "".join(f"<li>{smiles}</li>" for smiles in invalid_smiles)
        html_parts.append(
            "<div style='color:#b00020;margin-top:12px;'>"
            "<strong>Invalid SMILES strings:</strong>"
            f"<ul style='margin:4px 0 0 16px;'>{invalid_list_items}</ul>"
            "</div>"
        )

    return "".join(html_parts)

# --- Gradio Interface ---
with gr.Blocks() as demo:
    gr.Markdown("## Molecule Renderer")
    gr.Markdown("Enter one or more SMILES strings to visualize the molecules.")

    with gr.Row():
        smiles_input = gr.DataFrame(
            value=[[smiles] for smiles in list_smiles],
            headers=["SMILES"],
            label="SMILES Strings",
            row_count=(len(list_smiles), "dynamic"),
            col_count=1,
            type="array",
            wrap=True,
        )
        render_image_button = gr.Button("Render Image")
        render_svg_button = gr.Button("Render SVG")

    gr.Markdown("#### Image")
    image_output = gr.Gallery(label="Molecule Structures")

    gr.Markdown("#### SVG")
    svg_output = gr.HTML(label="SVG Output")

    # Separate actions for independent rendering
    render_image_button.click(
        fn=render_images,
        inputs=smiles_input,
        outputs=image_output,
    )

    render_svg_button.click(
        fn=render_svgs,
        inputs=smiles_input,
        outputs=svg_output,
    )

if __name__ == "__main__":
    demo.launch()