#!/usr/bin/env python3
# (C) 2026 Cadence Design Systems, Inc. (Cadence)
# All rights reserved.
# TERMS FOR USE OF SAMPLE CODE The software below ("Sample Code") is
# provided to current licensees or subscribers of Cadence products or
# SaaS offerings (each a "Customer").
# Customer is hereby permitted to use, copy, and modify the Sample Code,
# subject to these terms. Cadence claims no rights to Customer's
# modifications. Modification of Sample Code is at Customer's sole and
# exclusive risk. Sample Code may require Customer to have a then
# current license or subscription to the applicable Cadence offering.
# THE SAMPLE CODE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED.  OPENEYE DISCLAIMS ALL WARRANTIES, INCLUDING, BUT
# NOT LIMITED TO, WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
# PARTICULAR PURPOSE AND NONINFRINGEMENT. In no event shall Cadence be
# liable for any damages or liability in connection with the Sample Code
# or its use.


"""Converts a molecule structure into an image file (supporting "jpg", "tiff", "gif")."""

import io
import os
import pathlib
import sys

from openeye import oechem, oedepict
from PIL import Image

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Depict molecule (supported image formats: svg, png, jpg, tiff, gif)."
__SCRIPT_TOOLKITS__ = ["oechem", "oedepict"]
__SCRIPT_CATEGORIES__ = ["depiction"]


def main() -> int:
    """Molecule depiction."""
    argv = sys.argv
    itf = oechem.OEInterface()
    oechem.OEConfigure(itf, InterfaceData)
    oedepict.OEConfigureImageOptions(itf)
    oedepict.OEConfigure2DMolDisplayOptions(itf)

    if not oechem.OEParseCommandLine(itf, argv):
        return os.EX_USAGE

    mol_filename = itf.GetString("--mol")
    image_filename: str | None = (
        itf.GetString("--image") if itf.HasString("--image") else None
    )

    _check_image_file(image_filename)

    ifs = oechem.oemolistream()
    if not ifs.open(mol_filename):
        oechem.OEThrow.Fatal("Cannot open input file!")

    mol = oechem.OEGraphMol()
    if not oechem.OEReadMolecule(ifs, mol):
        oechem.OEThrow.Fatal("Cannot read input file!")

    width, height = oedepict.OEGetImageWidth(itf), oedepict.OEGetImageHeight(itf)
    opts = oedepict.OE2DMolDisplayOptions(width, height, oedepict.OEScale_AutoScale)
    oedepict.OESetup2DMolDisplayOptions(opts, itf)

    oedepict.OEPrepareDepiction(mol)
    render_molecule(mol, opts, image_filename)

    return os.EX_OK


def render_molecule(
    mol: oechem.OEMolBase,
    opts: oedepict.OE2DMolDisplayOptions,
    image_filename: str | None,
) -> None:
    """Render molecule to image."""
    disp = oedepict.OE2DMolDisplay(mol, opts)

    if image_filename is None:
        _image = oedepict.OEImage(disp.GetWidth(), disp.GetHeight())
        oedepict.OERenderMolecule(_image, disp)
        _img = Image.open(io.BytesIO(oedepict.OEWriteImageToBytes("png", _image)))
        _img.show()
        return

    ext = oechem.OEGetFileExtension(image_filename).lower()
    if oedepict.OEIsRegisteredImageFile(ext):
        ofs = oechem.oeofstream()
        if not ofs.open(image_filename):
            oechem.OEThrow.Fatal(f"Cannot open output file '{image_filename}'!")
        oedepict.OERenderMolecule(ofs, ext, disp)

    elif ext in ["tiff", "gif", "jpg"]:
        image_bytes: bytes = oedepict.OERenderMoleculeToBytes("png", disp)
        imagefile = io.BytesIO(image_bytes)
        image = Image.open(imagefile)

        if ext == "jpg":
            rgb_image = image.convert("RGB")
            rgb_image.save(image_filename, "JPEG")
        else:
            image.save(image_filename)


def _check_image_file(image_filename: str | None) -> None:
    # script will terminate if there is some issues
    if not image_filename:
        # image will be displayed on the screen as png
        return
    ext = pathlib.Path(image_filename).suffix[1:].lower()
    if not oedepict.OEIsRegisteredImageFile(ext) and ext not in ["jpg", "tiff", "gif"]:
        oechem.OEThrow.Fatal("Unknown image output type!")

    ofs = oechem.oeofstream()
    if not ofs.open(image_filename):
        oechem.OEThrow.Fatal("Cannot open output image file!")


InterfaceData = """
!CATEGORY "input/output options"

    !PARAMETER --mol
      !ALIAS -m
      !TYPE string
      !REQUIRED true
      !KEYLESS 1
      !VISIBILITY simple
      !BRIEF Input molecule file
    !END

    !PARAMETER --image
      !ALIAS -i
      !TYPE string
      !REQUIRED false
      !KEYLESS 2
      !VISIBILITY simple
      !BRIEF Output image file (PNG, SVG, PDF, "JPG", "TIFF", "GIF").
    !END

!END
"""


setattr(main, "__SCRIPT_NAME__", __SCRIPT_NAME__)
setattr(main, "__SCRIPT_DESC__", __SCRIPT_DESC__)
setattr(main, "__SCRIPT_TOOLKITS__", __SCRIPT_TOOLKITS__)
setattr(main, "__SCRIPT_CATEGORIES__", __SCRIPT_CATEGORIES__)

if __name__ == "__main__":
    sys.exit(main())