#!/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.


"""Code snippet for extracting compounds from a file based on molecule title."""

import argparse
import os
import sys
from pathlib import Path

from openeye import oechem
from rich.console import Console


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(
        description="Extract compounds from a file based on molecule title.",
    )
    parser.add_argument("-i", "--input", required=True, help="Input file name")
    parser.add_argument("-o", "--output", required=True, help="Output file name")

    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("-t", "--title", help="Single mol title to extract")
    group.add_argument(
        "-l", "--list", type=Path, help="List file of mol titles to extract"
    )

    return parser.parse_args()


def main() -> int:
    """Extract compounds from a file based on molecule title."""
    args = parse_args()
    console = Console()

    ifs = oechem.oemolistream()
    if not ifs.open(args.input):
        oechem.OEThrow.Fatal(f"Unable to open {args.input} for reading")

    ofs = oechem.oemolostream()
    if not ofs.open(args.output):
        oechem.OEThrow.Fatal(f"Unable to open {args.output} for writing")

    title_set: set[str] = set()
    if args.list:
        list_path: Path = args.list
        if not list_path.exists():
            oechem.OEThrow.Fatal(f"Unable to open {list_path} for reading")
        for name in list_path.read_text().splitlines():
            title = name.strip()
            if title:
                title_set.add(title)
    elif args.title:
        title_set.add(args.title)

    if len(title_set) == 0:
        oechem.OEThrow.Fatal("No titles requested")

    mol_extract(ifs, ofs, title_set, console=console)
    console.print(
        f"Extracted {len(title_set)} title(s) from {Path(args.input).name} to {Path(args.output).name}"
    )

    return os.EX_OK


def mol_extract(
    ifs: oechem.oemolistream,
    ofs: oechem.oemolostream,
    title_set: set[str],
    console: Console,
) -> None:
    """Extract molecules from a database whose titles match the given set."""
    mol_database = oechem.OEMolDatabase(ifs)

    for idx in range(mol_database.GetMaxMolIdx()):
        title = mol_database.GetTitle(idx)
        if title in title_set:
            console.print(f"Extracting {title} (mol index {idx+1})")
            mol_database.WriteMolecule(ofs, idx)


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