"""
enamine_functions.py - functions for interacting with Enamine Web Services

Use of this code is subject to the terms of the StarDrop License Agreement and
use, copying or redistribution of this file for use by other persons or
organisations is prohibited

Copyright(C) Optibrium Ltd 2026
support@optibrium.com
"""
import time

import PythonWidgets
import QueryInterfaceWidgets
import enamine_query_logger
from enamine_exceptions import (
    EnamineException, DatasetCreationException, DelayedResultNotFound, ConnectionException, AuthenticationException
)
from enamine_web_services import EnamineWebServices
import enamine_credentials
import app
import socket

from enamine_data_utils import convert_json_to_data_list

STRUCTURE_QUERY_WAIT = 1.0
INVALID_CREDENTIALS_ERROR = "Invalid credentials. Please try again."

log = enamine_query_logger.EnamineQueryLogger('enaminequerytool.log')
socket.setdefaulttimeout(20)

cred = None


def _get_credentials(appui):
    """Get a username and password from the user and store in a CredentialsManager object for use during the lifetime
    of the session.
    """
    global cred
    username = ""
    password = ""
    if not cred or cred.configured is False:
        dlg = PythonWidgets.LogInDialog()
        dlg.setWindowTitle("Enamine Login")
        if dlg.run():
            try:
                username = dlg.getUser()
                password = dlg.getPassword()
                if username == "" or password == "":
                    appui.show(
                        "User and password must be completed")
                else:
                    cred = enamine_credentials.CredentialsManager(username, password)
            except ValueError as ve:
                print(ve)
            except Exception as e:
                print(e)
        else:
            return  # ignore if cancelled
    else:

        credentials = vars(cred)
        username = credentials.get('_username')
        password = credentials.get('_password')

    if not username or not password:
        raise ValueError("Could not locate authentication details for Enamine "
                         "API - cannot create link to web services")

    return username, password


def _clear_credentials():
    """ clears cached credentials to allow signing in again"""
    global cred
    cred = None


def _create_enamine_web_services(appui):
    """
    Create an instance of the EnamineWebServices object either based on user
    input for details or from the session credentials.

    If the credentials have not been confirmed or do not exist then the user will be
    prompted to provide new details.

    This function may throw ValueError if after all this logic it still has no
    username or password.

    :return: an instance of EnamineWebServices
    :rtype: EnamineWebServices
    """
    username, password = _get_credentials(appui)

    if not username or not password:
        raise ValueError("Could not locate authentication details for Enamine "
                         "API - cannot create link to web services")

    web_services = EnamineWebServices(username, password)
    web_services.login()
    return web_services


def _create_string_entry(value):
    """

    :param value:   The string value to store in the new string entry
    :return:    The new string data entry
    """
    return app.Dataset.string(str(value))


def _create_columns(column_definitions):
    """
    Create columns for a dataset based on a set of column definitions. This
    is a convenience function that allows us to cut down on lots of repeated
    code to make columns.

    :param column_definitions:  The list of columns to create. Each column
                                should consist of a name, a type and None or a
                                list of category names depending on type
    :return: A tuple containing column headers and category dictionaries
    """
    cols = []

    for column_definition in column_definitions:
        header, col_type, _ = column_definition

        if col_type == app.mol:
            cols.append(app.Dataset.meta(header, app.mol))
        elif col_type == app.cont:
            cols.append(app.Dataset.createNumberHeader(header))
        elif col_type == app.string:
            cols.append(app.Dataset.meta(header))

    return cols


def find_threshold(appui):
    """
    Function for requesting a similarity threshold percentage from the user.

    :param appui:       The StarDrop appui object
    :return threshold:  The similarity threshold input by the user as a float
    """
    threshold = (appui.query('Similarity',
                             'Enter a similarity value:\n'
                             '(Note: This must be 0.4 or higher. If the value you choose is low, you may\n'
                             'see errors because too many results are returned)',
                             '0.8'))

    threshold = float(threshold)

    # The threshold needs to be 0.4 or above, otherwise the API won't return results.
    if threshold < 0.40 or threshold > 1:
        threshold = find_threshold(appui)

    return threshold


def sketch_structure_query(appui):
    """
        Perform a structural search query based on the provided structure and using
        the specified search type (and similarity range if specified).

        :param appui:    The StarDrop appui object
        :return: None
        """
    try:
        smiles = ""
        operator = ""
        similarity = 0

        web_services = _create_enamine_web_services(appui)

        dlg = QueryInterfaceWidgets.StructureQueryDialog()
        dlg.setPossibleMatchTypes(["SUBSTRUCTURE", "SIMILARITY", "EXACT MATCH"])
        if dlg.run():
            query = dlg.selectedQuery()
            smiles, operator, similarity = query[:3]
            log.log("running Structure query %s" % smiles)
            print("Running structure query for " + smiles)
        else:
            return  # ignore if cancelled

        run_structure_query(appui, web_services, smiles, operator, similarity)

    except Exception as e:
        log.log(repr(e))
        if isinstance(e, AuthenticationException):
            _clear_credentials()
            appui.show(INVALID_CREDENTIALS_ERROR)


def structure_query_selected(appui):
    """
        Perform a structural search query using the selected structure and
        an input similarity value.

        :param appui:    The StarDrop appui object
        :return: None
        """

    try:
        web_services = _create_enamine_web_services(appui)

        # get currently selected molecule
        smiles_input = appui.getCurrentMolecules()

        # if no molecule is selected, tell user to select one and re-run
        if not smiles_input:
            appui.show("Please select a molecule before searching for similar compounds")
        elif len(smiles_input) > 1:
            appui.show("Please select exactly one molecule")
        else:
            smiles_input = smiles_input[0]
            print("Running structure query for " + smiles_input)

            threshold = find_threshold(appui)
            threshold_value = str(threshold)

            operator = "LIKE"

            run_structure_query(appui, web_services, smiles_input, operator, threshold_value)

    except Exception as e:
        log.log(repr(e))
        if isinstance(e, AuthenticationException):
            _clear_credentials()
            appui.show(INVALID_CREDENTIALS_ERROR)


def run_structure_query(appui, web_services, smiles, operator, similarity):
    """
    Obtain the other parameters from the user, use the web services chemical_search function to get
    the data and when it is returned package it into a dataset.

    :param: appui               The StarDrop appui object
    :param: web_services        The web_services instance
    :param: smiles              The input SMILES string
    :param: operator            Whether the search should be EXACT, SIMILAR or SUBSTRUCTURE
    :param: similarity          The similarity value
    :return: dataset            The dataset of compounds found in Enamine database
    """
    try:
        cb = PythonWidgets.ComboBoxSelector()

        cb.setWindowTitle("Set up parameters")
        cb.setText(
            "Search mode options: \n"
            "Building blocks are scaffolds and intermediates \n"
            "Screening compounds are for biological screening")
        options = {
            "Search mode": [
                "Building blocks",
                "Screening compounds",
                "REAL database",
                "Made Building blocks",
                "Bioactive compounds"
            ],
            "Currency": [
                "USD", "EUR"
            ]
        }
        cb.setOptions("Search mode: ", options["Search mode"])
        cb.addOptions("Currency: ", options["Currency"])

        if cb.run():  # True if user clicked OK rather than Cancel
            catalog = cb.getSelectedOption(0)
            currency = cb.getSelectedOption(1)
        else:
            return  # ignore if cancelled

        try:
            print("Catalogue: " + catalog)

            prog_dlg = PythonWidgets.SimpleProgress()
            prog_dlg.modal = True
            prog_dlg.setWindowTitle('Find Enamine Compounds')
            prog_dlg.setMessage('Running Enamine Web Services search')
            prog_dlg.show()
            prog_dlg.setProgress(4)
            appui.updateUI()

            prog_dlg.setProgress(10)
            # If the user sketched the structure and chose a similarity lower than 0.4, tell them to set this
            # higher and bring up threshold dialog.
            if operator == "LIKE":
                if float(similarity) < 0.4:
                    appui.show("Please choose a higher similarity")
                    similarity = find_threshold(appui)

            prog_dlg.setProgress(15)

            try:
                search_uuid = web_services.chemical_search(smiles, operator, catalog, similarity, currency)

                time.sleep(STRUCTURE_QUERY_WAIT)

                molecules = web_services.fetch_delayed_result(search_uuid, prog_dlg)
            except ConnectionException as ex:
                log.log(repr(ex))
                print(ex)
                appui.show("Error running Enamine Web Services search. Please try again later.")
                return
            except DelayedResultNotFound as ex:
                log.log(repr(ex))
                print(ex)
                appui.show("Timed out waiting for search results from Enamine Web Services. Please try again later.")
                return

            bin_dataset = _export_enamine_dataset(molecules, prog_dlg, currency=currency)
            appui.exportDS("Structure Query Result", bin_dataset)

            prog_dlg.setProgress(90)

            prog_dlg.hide()
            prog_dlg.delete()

        except DatasetCreationException as e:
            print(e)
            appui.show("No match found: " + str(e))
        except Exception as e:
            print(e)
            log.log(repr(e))

    except Exception as e:
        log.log(repr(e))
        print(e)


def selected_ids_query(appui):
    """
    Perform an ID query with ID(s) entered into a dialog box.

    One or multiple IDs can be given.

    :param appui:    The StarDrop appui object
    :return: None
    """
    try:
        ids = []
        currency = ""

        web_services = _create_enamine_web_services(appui)

        dlg = QueryInterfaceWidgets.IDQueryDialog()

        if dlg.run():
            # The return value is simply the list of IDs
            ids = dlg.selectedQuery()
            log.log("running ID query %s" % (repr(ids)))

            cb = PythonWidgets.ComboBoxSelector()
            cb.setWindowTitle("Set up parameters")
            cb.setText("")
            options = {
                "Currency": [
                    "USD", "EUR"
                ]
            }
            cb.setOptions("Currency: ", options["Currency"])
            if cb.run():  # True if user clicked OK rather than Cancel
                currency = cb.getSelectedOption(0)
            else:
                return  # ignore if cancelled
        else:
            return  # ignore if cancelled

        prog_dlg = PythonWidgets.SimpleProgress()
        prog_dlg.setWindowTitle('Find Enamine Compounds')
        prog_dlg.setMessage('Running Enamine Web Services search')
        prog_dlg.show()
        prog_dlg.setProgress(5)

        print(ids)

        # If the user has entered IDs separated by commas or spaces, ids will be a tuple of length one with
        # IDs separated by a comma or space so repackage these into a tuple
        for mol_id in ids:
            if len(ids) == 1:
                if ", " in mol_id:
                    ids = mol_id.split(", ")
                elif " " in mol_id:
                    ids = mol_id.split(" ")

            # If the user has entered multiple IDs on different lines AND separated by a comma, ask them
            # to enter one per line
            if len(ids) > 1 and ", " in mol_id:
                appui.show("Please write one ID per line")
                dlg = QueryInterfaceWidgets.IDQueryDialog()
                if dlg.run():
                    # The return value is simply the list of IDs
                    ids = dlg.selectedQuery()
                    log.log("running ID query %s" % (repr(ids)))
                else:
                    return  # ignore if cancelled

        if len(ids) > 0:
            try:

                prog_dlg.setProgress(10)

                try:
                    molecule_data = web_services.load_molecules(ids, currency)
                except EnamineException as ex:
                    appui.show(str(ex))
                    log.log(repr(ex))
                    return

                bin_dataset = _export_enamine_dataset(molecule_data, prog_dlg, currency)
                appui.exportDS("ID Query Result", bin_dataset)

            except Exception as e:
                log.log(repr(e))
                print(e)
                if isinstance(e, DatasetCreationException):
                    appui.show("No match found: " + str(e))
    except Exception as e:
        log.log(repr(e))
        if isinstance(e, AuthenticationException):
            _clear_credentials()
            appui.show(INVALID_CREDENTIALS_ERROR)


def _export_enamine_dataset(molecule_data, prog_dlg, currency: str = 'USD'):
    data_list = convert_json_to_data_list(molecule_data)

    prog_dlg.setProgress(60)

    column_definitions = [
        ("SMILES", app.mol, None),
        ("Enamine ID", app.string, None),
        ("Formula", app.string, None),
        ("Molecular Weight", app.cont, None),
        ("CAS", app.string, None),
        ("Catalogue", app.string, None),
        ("Amount in stock (g)", app.string, None),
        ("Price (" + str(currency) + ")", app.cont, None),
    ]

    cols = _create_columns(column_definitions)

    prog_dlg.setProgress(70)

    dataset = app.Dataset(cols)
    dataset.idColName = "Id"

    prog_dlg.setProgress(75)

    for row in data_list:
        dataset.cols[0].append(
            app.Dataset.molecule(row[0]))
        dataset.cols[1].append(
            _create_string_entry(row[1]))
        dataset.cols[2].append(
            _create_string_entry(row[2]))
        dataset.cols[3].append(
            app.Dataset.cont(row[3], 0.0))
        dataset.cols[4].append(
            _create_string_entry(row[4]))
        dataset.cols[5].append(
            _create_string_entry(row[5]))
        dataset.cols[6].append(
            _create_string_entry(row[6]))
        dataset.cols[7].append(
            app.Dataset.cont(row[7], 0.0))

    prog_dlg.setProgress(95)

    if dataset.cols[0]:
        pack_data = None
        try:
            pack_data = dataset.pack()
        except Exception as e:
            print("type error: " + str(e))
        return pack_data

    raise DatasetCreationException


def get_stardrop_definitions():
    definitions = [
        {
            'script_name': 'Enamine/Sketch structure to use in search...',
            'callback': sketch_structure_query
        },
        {
            'script_name': 'Enamine/Search by similarity to selected molecule...',
            'callback': structure_query_selected
        },
        {
            'script_name': 'Enamine/Search by Enamine ID...',
            'callback': selected_ids_query
        }
    ]

    return definitions
