"""
'Query Round Trip' example script
Usage:
    python <script_name>.py <molecule_file_containing_query>
"""
from orionclient.session import APISession
from orionclient.types import MolsearchQuery, MolsearchDatabase
from openeye.oechem import OEMol, oemolistream, OEReadMolecule, OEThrow
import time
import sys


def monitor_query_until_complete(query):
    state = ''

    # Wait for completion
    while query.state != 'SUCCESS':
        APISession.refresh_resource(query)
        if state != query.state:
            print(f"Query state: {query.state}")
        time.sleep(1)

        state = query.state

    if state == '':
        print(f"Query state: {query.state}")


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


if __name__ == '__main__':
    PROJECT = APISession.get_current_project().id
    USER = APISession.get_user_profile()

    db = find_first_available_db()

    ifs = oemolistream()
    filename = sys.argv[1]
    if not ifs.open(filename):
        OEThrow.Fatal(f"Unable to open file {filename}")

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

    # Create molecule search query
    subsearch_query = MolsearchQuery.create_subsearch_query(
        database_id=db.id,
        num_hits=100,
        mdlquery=mol,
        subsearch_query_type="MDLJSON",
        name=f"{filename} subsearch",
        project=PROJECT,
        aliphatic_constraint=False,
        topology_constraint=False,
        stereo_constraint=False,
        isotope_constraint=False,
        session=APISession,
    )

    if subsearch_query.state in ['QUEUED', 'PROCESSING', 'SUCCESS']:
        monitor_query_until_complete(subsearch_query)
    else:
        raise SystemExit("Query was found to be in an unexpected state,"
                         f" {subsearch_query.state}")

    # Export query results to an Orion dataset
    ds_id = subsearch_query.export(
        name=subsearch_query.name,
        session=APISession,
        project=PROJECT,
        path=f"/project/{PROJECT}/user/{USER.username}",
    )

    print(f"Exported molecule search query to dataset with ID {ds_id}")

    # Delete molecule search query
    APISession.delete_resource(subsearch_query)

    print("Deleted the associated query")