import gradio as gr
from typing import List
from utils import DbProcessor
import os

db_processor = DbProcessor()

def get_data(
        smiles: str, iupac: str, name: str, inchl: str, inchl_key: str, return_fields: List[str]
        ):
    smiles = smiles.strip()
    iupac = iupac.strip()
    name = name.strip()
    inchl = inchl.strip()
    inchl_key = inchl_key.strip()

    inputs = {
        "SMILES": smiles,
        "IUPAC": iupac,
        "NAME": name,
        "InChI": inchl,
        "InChIKey": inchl_key
    }
    
    count = 0
    query_type = None
    for k, v in inputs.items():
        if v:
            count += 1
            query_type = k

    if count == 0 or count > 1:
        return [f"For search, you should specify only one field, but you specified {count} fields"] * 5

    result_template = {
        'SMILES': "",
        'IUPAC': "",
        'InChI': "",
        'InChIKey': "", 
        'synonyms': ""
    }

    try:
        result = db_processor.request_data(text_input=inputs[query_type], query_type=query_type, return_fields=return_fields)
        for key, value in result.items():
            if key in result_template:
                result_template[key] = value
        if "synonyms" in return_fields:
            result_template["synonyms"] = str(result["synonyms"])[1:-1]

        return [result_template['SMILES'], result_template['IUPAC'], result_template['synonyms'], result_template['InChI'], result_template['InChIKey']]
    
    except Exception as e:
        print(e)
        return ["Compound not found"] * 5

main_interface = gr.Interface(
    fn=get_data,
    allow_flagging='never',
    inputs=[
        gr.Textbox(label="Search from SMILES", placeholder="Enter SMILES name here"),
        gr.Textbox(label="Search from IUPAC", placeholder="Enter IUPAC name here"),
        gr.Textbox(label="Search from trivial name", placeholder="Enter name here"),
        gr.Textbox(label="Search from InChI", placeholder="Enter InChI here"),
        gr.Textbox(label="Search from InChIKey", placeholder="Enter InChIKey here"),
        gr.CheckboxGroup(choices=["IUPAC", "SMILES", "InChI", "InChIKey", "synonyms"], label="Choose data fields to request")
    ],
    outputs=[
        gr.Textbox(label="SMILES name"), 
        gr.Textbox(label="IUPAC name"), 
        gr.Textbox(label="Synonyms"), 
        gr.Textbox(label="InChI"), 
        gr.Textbox(label="InChIKey")
    ],
    examples=[
        ["", "", "(+)-PSI reagent", "", "", ["IUPAC", "SMILES", "InChI", "InChIKey", "synonyms"]],
        ["CCCCO", "", "", "", "", ["IUPAC", "synonyms"]],
        ["C[C@@H](C1=CC=C(C=C1)CC(C)C)C(=O)O", "", "", "", "", ["IUPAC", "SMILES", "InChI", "InChIKey", "synonyms"]],
    ],
    theme=gr.themes.Base()
)

if __name__ == "__main__":
    main_interface.launch(share=True, auth=(os.getenv('user'), os.getenv('password')))