"""
MolPortFunctions.py - functions for interacting with MolPort

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 PythonWidgets
import QueryInterfaceWidgets

from script_logger import ScriptLogger
from MolPortWebServices import MolPortWebServices, SearchTypes
import json
import os
import app
import socket

MOLECULAR_WEIGHT = "Molecular Weight"
BUILDING_BLOCK_SUPPLIERS = "Building Block Suppliers"
VERIFIED_AMOUNT = "Verified Amount"
LARGEST_STOCK = "Largest Stock"
MOLPORT_ID = "MolPort ID"
SCREENING_BLOCK_SUPPLIERS = "Screening Block Suppliers"
VIRTUAL_SUPPLIERS = "Virtual Suppliers"
CANONICAL_SMILES = "Canonical SMILES"
UNVERIFIED_AMOUNT = "Unverified Amount"
MOLPORT_SHOP_URL = "https://www.molport.com/shop/compound"

log = ScriptLogger('molport.log')
# TODO: Rethink when scripting is not in GUI thread
socket.setdefaulttimeout(20)

try:
    import refreshscript


    def refresh_molport_search(method, params):
        """
        Run a search and supply the information to StarDrop to perform a refresh operation.
        :param method: Should be 'MolPort'.
        :param params: The parameters to use when running the search.
        """
        if method == "MolPort":
            print('refreshing Molport')
            ds = runSearch(None,
                           int(params.get("search", "")),
                           int(params.get("vault", "")),
                           params.get("projects", "").split(','),
                           params.get("token", ""))
        return ds.pack()


    refreshscript.setRefreshFunction(refresh_molport_search)
except Exception:
    pass


def _serialise_auth_details(username: str, password: str, confirmed: bool = False):
    """
    Serialise the provided authentication details to a molport_auth.txt file.
    The confirmed parameter determines whether there has been a successful run
    of the web services with these details yet. If yes then there is no further
    need to keep prompting for the details, however if there hasn't yet been
    one then we should continue to query until we get one which works.

    :param username:    The username to store in the auth file
    :param password:    The authentication code to store in the auth file
    :param confirmed:   Whether these details have been used successfully
                        or not. Default = False
    :return: None
    """
    auth_file_path = "molport_auth.txt"

    auth_details = {
        "Username": username,
        "Password": password,
        "Confirmed": confirmed
    }

    auth_string = json.dumps(auth_details)

    try:
        with open(auth_file_path, "w") as f:
            f.write(auth_string)
    except EnvironmentError:
        log.log("Failed to write authentication details to file")


def _request_login_details():
    """
    Query the user for a login and password combination we can use as details
    with the MolPort API

    :return: A tuple formed from the username and password respectively
    """
    dlg = PythonWidgets.LogInDialog()
    if dlg.run():
        username = dlg.getUser()
        password = dlg.getPassword()

        _serialise_auth_details(username, password)

        return username, password
    else:
        return None, None


def _create_mol_port_web_services() -> MolPortWebServices:
    """
    Create an instance of the MolPortWebServices object either based on user
    input for details or from the stored molport_auth.txt depending on whether
    stored details exist, and if they do whether they have been confirmed or
    not.

    If they have not been confirmed or do not exist then the user will be
    prompted to provide new details, otherwise we proceed with those in the
    stored file.

    This function may throw ValueError if after all this logic it still has no
    username or password - for example if there was no stored file and the user
    then cancelled the LogInDialog which popped up instead of providing details.

    :return: an instance of MolPortWebServices
    """
    username = None
    password = None
    confirmed = False
    auth_file_path = "molport_auth.txt"

    if not os.path.isfile(auth_file_path):
        username, password = _request_login_details()
    else:
        try:
            with open(auth_file_path, "r") as f:
                file_contents = f.read()
                auth_details = json.loads(file_contents)
                if "Username" in auth_details:
                    username = auth_details["Username"]
                if "Password" in auth_details:
                    password = auth_details["Password"]
                if "Confirmed" in auth_details:
                    confirmed = auth_details["Confirmed"]

                if not confirmed:
                    username, password = _request_login_details()

        except ValueError:
            log.log(
                "Failed to read %s as JSON - invalid format" % auth_file_path)
        except EnvironmentError:
            log.log("Failed to open %s" % auth_file_path)

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

    web_services = MolPortWebServices(username, password)
    return web_services


def _convert_structure_operator_to_search_type(structure_operator):
    """
    Convert a StarDrop structure operator (such as "IN" or "LIKE") into an
    equivalent MolPort API SearchType value. We currently also return a value
    for "NOTIN" despite the fact there is no sensible mapping for this. In the
    future it would maek a lot of sense if this was removed.

    :param structure_operator:  The structure operator to convert
    :return: A SearchType corresponding to the operator provided
    :rtype: int
    """
    if structure_operator == "IN":
        return SearchTypes.SUBSTRUCTURE
    elif structure_operator == "NOTIN":
        # TODO: Resolve the fact that there is no negated search in MolPort
        return SearchTypes.SUBSTRUCTURE
    elif structure_operator == "LIKE":
        return SearchTypes.SIMILARITY
    elif structure_operator == "EQ":
        return SearchTypes.EXACT


def _create_cont_entry(value, error):
    """
    Create a Cont entry - this function was taken almost entirely from
    DatasetCreator.py in the QueryTool

    :param value:   The value to store in the new cont entry
    :param error:   The error to store in the new cont entry
    :return:    A tuple representing a new cont entry which can be appended to a
                column
    """
    qual = ''
    try:
        val = float(value)
    except ValueError:
        if value[0] in ('<', '>'):
            qual = value[0]
        try:
            val = float(value[1:])
        except ValueError:
            val = float(value[2:])
    return app.Dataset.cont(val, error, qualifier=qual)


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(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 = []
    category_dicts = []

    for column_definition in column_definitions:
        header, col_type, categories = 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))
        elif col_type == app.cat:
            categories = set(categories)
            categories.discard('')
            header, cat_dict = app.Dataset.createCategoryHeader(header,
                                                                categories)
            cols.append(header)
            category_dicts.append(cat_dict)

    return cols, category_dicts


def _parse_supplier(dataset, molecule, suppliers, supplier_string,
                    category_dict):
    """
    Parse the provided supplier object to create dataset rows from it

    :param dataset:         The dataset to write new rows to
    :param molecule:        The set of molecule data
    :param suppliers:       The suppliers object to represent
    :param supplier_string: The category string to use for this supplier type
    :param category_dict:   A dictionary of categories to use as necessar
    :return: None
    """
    mol_id = molecule["Id"]
    id_string = f"<a href=\"{MOLPORT_SHOP_URL}/{mol_id}\">{mol_id}</a>"

    for supplier in suppliers:
        for costs in supplier.get("Shipment Costs", []):
            dataset.cols[0].append(
                app.Dataset.molecule(molecule["SMILES"]))
            dataset.cols[1].append(
                _create_string_entry(molecule["Molport Id"]))
            dataset.cols[2].append(
                _create_string_entry(id_string))
            dataset.cols[3].append(
                _create_string_entry(molecule["Formula"]))
            dataset.cols[4].append(
                _create_cont_entry(molecule[MOLECULAR_WEIGHT], 0.0))
            dataset.cols[5].append(
                _create_string_entry(molecule[LARGEST_STOCK]))
            dataset.cols[6].append(
                app.Dataset.createCategoryEntry(supplier_string, category_dict))
            dataset.cols[7].append(
                _create_cont_entry(costs["Price"], 0.0))
            dataset.cols[8].append(
                _create_string_entry(costs["Currency"]))
            dataset.cols[9].append(
                app.Dataset.createInvalidEntry(app.cont))
            dataset.cols[10].append(
                app.Dataset.createInvalidEntry(app.string))
            dataset.cols[11].append(
                _create_string_entry(costs["Location"]))
            dataset.cols[12].append(
                app.Dataset.createInvalidEntry(app.string))
            dataset.cols[13].append(
                app.Dataset.createInvalidEntry(app.string))

        for catalog in supplier.get("Catalogues", []):
            for packing in catalog.get("Available Packings", []):
                dataset.cols[0].append(
                    app.Dataset.molecule(molecule["SMILES"]))
                dataset.cols[1].append(
                    _create_string_entry(molecule["Molport Id"]))
                dataset.cols[2].append(
                    _create_string_entry(id_string))
                dataset.cols[3].append(
                    _create_string_entry(molecule["Formula"]))
                dataset.cols[4].append(
                    _create_cont_entry(molecule[MOLECULAR_WEIGHT], 0.0))
                dataset.cols[5].append(
                    _create_string_entry(molecule[LARGEST_STOCK]))
                dataset.cols[6].append(
                    app.Dataset.createCategoryEntry(BUILDING_BLOCK_SUPPLIERS, category_dict))
                dataset.cols[7].append(
                    _create_cont_entry(packing["Price"], 0.0))
                dataset.cols[8].append(
                    _create_string_entry(packing["Currency"]))
                dataset.cols[9].append(
                    _create_cont_entry(packing["Amount"], 0.0))
                dataset.cols[10].append(
                    _create_string_entry(packing["Measure"]))
                dataset.cols[11].append(
                    app.Dataset.createInvalidEntry(app.string))
                dataset.cols[12].append(
                    _create_string_entry(catalog["Catalog Number"]))
                dataset.cols[13].append(
                    app.Dataset.string(
                        str(catalog["Last Update Date"])))


def _create_id_dataset(molecule_data):
    """
    Create ID dataset for the provided molecule data

    There are currently assumptions about a fixed set of columns and a fixed
    format returned by the API so changes will require code tweaks here.

    :param molecule_data:   The molecule data to parse into a dataset
    :return: a dataset object for the provided data or None if there wasn't any
    """
    if not molecule_data:
        print("No Molecule data - exiting")
        return None

    column_definitions = [
        ("SMILES", app.mol, None),
        (MOLPORT_ID, app.string, None),
        ("ID", app.string, None),
        ("Formula", app.string, None),
        (MOLECULAR_WEIGHT, app.cont, None),
        (LARGEST_STOCK, app.string, None),
        ("Catalogue Type", app.cat,
         [SCREENING_BLOCK_SUPPLIERS, BUILDING_BLOCK_SUPPLIERS, VIRTUAL_SUPPLIERS]),
        ("Price", app.cont, None),
        ("Currency", app.string, None),
        ("Amount", app.cont, None),
        ("Measure", app.string, None),
        ("Location", app.string, None),
        ("Catalog Number", app.string, None),
        ("Last Update Date", app.string, None),
    ]

    cols, category_dicts = _create_columns(column_definitions)
    category_dict = category_dicts[0]

    dataset = app.Dataset(cols)
    dataset.idColName = MOLPORT_ID

    for molecule in molecule_data:
        screening_suppliers = molecule["Catalogues"][
            SCREENING_BLOCK_SUPPLIERS]
        building_suppliers = molecule["Catalogues"][BUILDING_BLOCK_SUPPLIERS]
        virtual_suppliers = molecule["Catalogues"][VIRTUAL_SUPPLIERS]

        if len(screening_suppliers) > 0:
            _parse_supplier(dataset, molecule, screening_suppliers,
                            SCREENING_BLOCK_SUPPLIERS, category_dict)
        if len(building_suppliers) > 0:
            _parse_supplier(dataset, molecule, building_suppliers,
                            BUILDING_BLOCK_SUPPLIERS, category_dict)
        if len(virtual_suppliers) > 0:
            _parse_supplier(dataset, molecule, virtual_suppliers,
                            VIRTUAL_SUPPLIERS, category_dict)

    return dataset


def id_query(appui):
    """
    Perform an ID query for a set of ID(s) specified by the user.

    :param appui:    The StarDrop appui object
    :return: None
    """
    try:
        web_services = _create_mol_port_web_services()

        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)))
            if len(ids) > 0:
                try:
                    data_list = []
                    found_data = False

                    for mol_id in ids:
                        molecule_data = web_services.load_molecule(appui, mol_id)

                        if molecule_data and "Data" in molecule_data:
                            data = molecule_data["Data"]
                            if "Molecule" in data:
                                data_list.append(data["Molecule"])
                                found_data = True

                    if not found_data:
                        warning_dlg = PythonWidgets.Warning()
                        warning_dlg.setMessage("Could not find molecules by ID")
                        warning_dlg.run()

                    dataset = _create_id_dataset(data_list)
                    if dataset:
                        appui.exportDS("IDQueryResult", dataset.pack())
                        _serialise_auth_details(web_services.username,
                                                web_services.authentication_code,
                                                True)
                except Exception as e:
                    log.log(repr(e))
            else:
                warning_dlg = PythonWidgets.Warning()
                warning_dlg.setMessage(
                    "No IDs specified, could not perform query")
                warning_dlg.run()
    except Exception as e:
        log.log(repr(e))


def selected_ids_query(appui):
    """
    Perform an ID query with ID(s) extracted from the ID field of selected
    rows.

    :param appui:    The StarDrop appui object
    :return: None
    """
    try:
        web_services = _create_mol_port_web_services()

        ids = appui.getSelectedIDs()
        log.log("running ID query %s" % (repr(ids)))
        if len(ids) > 0:
            try:
                data_list = []
                found_data = False

                prog_dlg = PythonWidgets.SimpleProgress()
                prog_dlg.modal = True
                prog_dlg.setWindowTitle('Find MolPort Compounds')
                prog_dlg.setMessage('Running MolPort search web service')
                prog_dlg.setMaximum(len(ids))
                prog_dlg.setProgress(0)
                prog_dlg.show()
                appui.updateUI()
                for i, mol_id in enumerate(ids):
                    molecule_data = web_services.load_molecule(appui, mol_id)

                    if molecule_data and "Data" in molecule_data:
                        data = molecule_data["Data"]
                        if "Molecule" in data:
                            data_list.append(data["Molecule"])
                            found_data = True
                    prog_dlg.setProgress(i)
                    appui.updateUI()
                prog_dlg.hide()

                if not found_data:
                    warning_dlg = PythonWidgets.Warning()
                    warning_dlg.setMessage("Could not find molecules by ID")
                    warning_dlg.run()

                dataset = _create_id_dataset(data_list)
                if dataset:
                    appui.exportDS("IDQueryResult", dataset.pack())
                    _serialise_auth_details(web_services.username,
                                            web_services.authentication_code,
                                            True)
            except Exception as e:
                log.log(repr(e))
        else:
            warning_dlg = PythonWidgets.Warning()
            warning_dlg.setMessage("No rows selected, could not perform query")
            warning_dlg.run()
    except Exception as e:
        log.log(repr(e))


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

    We default to a maximum of 5 seconds per query in the API, and will try to
    return the maximum of 1000 rows if possible.

    :param appui:    The StarDrop appui object
    :return: None
    """
    try:
        web_services = _create_mol_port_web_services()

        dlg = QueryInterfaceWidgets.StructureQueryDialog()
        dlg.setSmartsEnabled(True)
        # Exclude "NO MATCH" from the set of possible match types
        dlg.setPossibleMatchTypes(["SUBSTRUCTURE", "SIMILARITY", "EXACT MATCH"])
        if dlg.run():
            query = dlg.selectedQuery()
            smiles, operator, similarity = query[:3]
            log.log("running Structure query %s" % smiles)
            try:
                search_type = _convert_structure_operator_to_search_type(
                    operator)
                prog_dlg = PythonWidgets.SimpleProgress()
                prog_dlg.modal = True
                prog_dlg.setWindowTitle('Find MolPort Compounds')
                prog_dlg.setMessage('Running MolPort search web service')
                prog_dlg.show()
                appui.updateUI()
                molecules = web_services.chemical_search(appui,
                                                         smiles,
                                                         search_type,
                                                         5000,
                                                         1000,
                                                         similarity)
                prog_dlg.hide()

                if "Data" in molecules:
                    if search_type == SearchTypes.SIMILARITY:
                        column_definitions = [
                            (CANONICAL_SMILES, app.mol, None),
                            (MOLPORT_ID, app.string, None),
                            ("ID", app.string, None),
                            (UNVERIFIED_AMOUNT, app.cont, None),
                            (VERIFIED_AMOUNT, app.cont, None),
                            ("Similarity Index", app.cont, None),
                        ]
                    else:
                        column_definitions = [
                            (CANONICAL_SMILES, app.mol, None),
                            (MOLPORT_ID, app.string, None),
                            ("ID", app.string, None),
                            (UNVERIFIED_AMOUNT, app.cont, None),
                            (VERIFIED_AMOUNT, app.cont, None),
                        ]
                    cols, _ = _create_columns(column_definitions)

                    dataset = app.Dataset(cols)
                    dataset.idColName = MOLPORT_ID
                    molecules = list(molecules["Data"]["Molecules"])
                    if len(molecules) > 0:
                        for molecule in molecules:
                            if "SMILES" in molecule:
                                if CANONICAL_SMILES in molecule:
                                    try:
                                        mol_id = molecule["Id"]
                                        id_string = f"<a href=\"{MOLPORT_SHOP_URL}/{mol_id}\">{mol_id}</a>"
                                        canonical_smiles_returned = molecule[CANONICAL_SMILES]
                                        dataset.cols[0].append(
                                            app.Dataset.molecule(canonical_smiles_returned))
                                        dataset.cols[1].append(
                                            _create_string_entry(molecule["MolPort Id"]))
                                        dataset.cols[2].append(
                                            _create_string_entry(id_string))
                                        dataset.cols[3].append(
                                            _create_cont_entry(molecule[UNVERIFIED_AMOUNT], 0.0))
                                        dataset.cols[4].append(
                                            _create_cont_entry(molecule[VERIFIED_AMOUNT], 0.0))
                                        if search_type == SearchTypes.SIMILARITY:
                                            dataset.cols[5].append(_create_cont_entry(
                                                str(1.0 - molecule["Similarity Index"]), 0.0))
                                    except Exception as e:
                                        print('Exception: %s' % repr(e))

                        if dataset:
                            appui.exportDS("StructureQueryResult",
                                           dataset.pack())
                            _serialise_auth_details(web_services.username,
                                                    web_services.authentication_code,
                                                    True)
                    else:
                        warning_dlg = PythonWidgets.Warning()
                        warning_dlg.setMessage(
                            "Could not find any matching molecules")
                        warning_dlg.run()
                        warning_dlg.delete()

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


def get_stardrop_definitions():
    definitions = [
        {
            'script_name': 'MolPort/By Structure...',
            'callback': structure_query
        },
        {
            'script_name': 'MolPort/By MolPort ID...',
            'callback': id_query
        },
        {
            'script_name': 'MolPort/Get Details For Selection',
            'callback': selected_ids_query
        }
    ]

    return definitions
