from orionclient.session import APISession
from orionclient.types import MolsearchQuery, MolsearchDatabase
from openeye.oechem import OEMol, oemolistream, OEReadMolecule, OEThrow


def find_first_available_db(search_type='2D'):
    if search_type.upper() not in ['2D', '3D']:
        raise ValueError("Search type must be '2D' or '3D'")

    search_filter = {'search_type': search_type.upper()}

    db = None
    for db in APISession.list_resources(MolsearchDatabase, search_filter):
        if db.state == "LOADED":
            break

    if db is None or db.state != "LOADED":
        raise ValueError(f"No loaded {search_type} database found")

    return db


PROJECT = APISession.get_current_project().id

db3d = find_first_available_db(search_type='3D')

ifs = oemolistream()
filename = "ibuprofen.sdf"

if not ifs.open(filename):
    OEThrow.Fatal(f"Unable to open file {filename}")

mol = OEMol()
OEReadMolecule(ifs, mol)
ifs.close()

fastrocs_query = MolsearchQuery.create_fastrocs_query(
    database_id=db3d.id,
    query_mol=mol,
    max_hits=100,
    name="Ibuprofen_3d_search",
    project=PROJECT,
    shape_only=False,
    sim_type="tanimoto",
    session=APISession,
)