import openai
import os
from openai import OpenAI
import gradio as gr 

api_key = os.getenv("OPENAI_API_KEY")

client = OpenAI(api_key=api_key)

cafchem_tools = [
		 {
        "type" : "mcp",
        "server_label":"CafChem_GenMask_gen_mask",
        "server_url":"https://cafierom-cafchem-genmask.hf.space/gradio_api/mcp/sse",
        "require_approval": "never",
        "allowed_tools": ["CafChem_GenMask_do_genmask"],
    },
    {
        "type" : "mcp",
        "server_label":"CafChem_Substitutions",
        "server_url":"https://cafierom-cafchem-substitutions.hf.space/gradio_api/mcp/sse",
        "require_approval": "never",
        "allowed_tools": ["CafChem_Substitutions_do_subrings"],
    },
		 {
        "type" : "mcp",
        "server_label":"CafChem_ADME_calc_adme",
        "server_url":"https://cafierom-cafchem-adme.hf.space/gradio_api/mcp/sse",
        "require_approval": "never",
        "allowed_tools": ["CafChem_ADME_calc_adme"],
    }
		 ]

gen_prompt = "Generate analogues of a hit molecule with the following SMILES string either \
by generative mask filling, masking HOWMASK of the tokens, \
or by making substitutions on phenyl rings, or by doing both, and report the QED value for each: INSERTSMILES "

med_prompt = "Would the molecule described by the following SMILES \
string be a feasible drug? INSERTSMILES"

gen_list = ["O=C1c3c(O/C(=C1/O)c2ccc(O)c(O)c2)cc(O)cc3O", "O=C(O)C[C@H](O)C[C@H](O)CCn2c(c(c(c2c1ccc(F)cc1)c3ccccc3)C(=O)Nc4ccccc4)C(C)C",
             "c1ccccc1C2=NCC(=O)N(C)c3ccc(Cl)cc23", "Oc1ccc(cc1)\C=C\c2cc(O)cc(O)c2"]

mols_list = ["CC(=O)Nc1ccc(O)cc1", "O=C(O)C[C@H](O)C[C@H](O)CCn2c(c(c(c2c1ccc(F)cc1)c3ccccc3)C(=O)Nc4ccccc4)C(C)C",
             "c1ccccc1C2=NCC(=O)N(C)c3ccc(Cl)cc23", "Oc1ccc(cc1)\C=C\c2cc(O)cc(O)c2"]

chat_history = []
global last_id
last_id = None

def chat(prompt, tools):
    chat_history.append(
                    {"role": "user", "content": prompt}
    )
    global last_id

    if tools == "Yes":
        if (last_id != None):
            response = client.responses.create(
                        model = "o4-mini",
                        tools = cafchem_tools,
                        input=prompt,
                        previous_response_id = last_id
            )
        else:
            response = client.responses.create(
                        model = "o4-mini",
                        tools = cafchem_tools,
                        input=prompt
            )
    else:
        if (last_id != None):
            response = client.responses.create(
                        model = "o4-mini",
        				input=prompt,
                        previous_response_id = last_id
            )
        else:
            response = client.responses.create(
                    model = "o4-mini",
    				input=prompt
            )

    chat_history.append(
                    {"role": "assistant", "content": response.output_text}
    )
    last_id = response.id
    
    return "", chat_history

def clear_history():
	global chat_history
	chat_history = []
	global last_id
	last_id = None

def add_mol(mol):
	new_prompt = med_prompt.replace("INSERTSMILES", mol)
	return new_prompt

def add_gen(mol,percent_mask):
	new_prompt = gen_prompt.replace("INSERTSMILES", mol).replace("HOWMASK",str(percent_mask))
	return new_prompt

with gr.Blocks() as forest:
  gr.Markdown(
      """
      # Chat with OpenAI 04-mini using the CafChem suite of tools.
      ### Enter your messages below including a SMILES string for a molecule, \
or choose a pre-set molecule from the list to auto-populate a question. You cans use the \
CafChem tools to generate analogues of a hit molecule for hit expansion or find some \
ADME properties of a molecule.
      """)

  tools = gr.Radio(choices = ["Yes", "No"],label="Use CafChem tools?",
                   interactive=True)
  
  chatbot = gr.Chatbot(type="messages")

  with gr.Row():
      msg = gr.Textbox(label="Type your messages here and hit enter.", scale = 2)
      chat_btn = gr.Button(value = "Send", scale = 0)

  gen_mols = gr.Radio(choices = gen_list,label="Generate new molecule hit suggestions: Quercetin, Lipitor, Diazepam, Resveratrol.",
                   interactive=True)
  percent_mask = gr.Radio(choices = [0.10, 0.15, 0.20],label="Fraction of hit molecule to mask.", value = 0.15,
                   interactive=True)
  mols = gr.Radio(choices = mols_list,label="ADME molecule suggestions: Paracetamol, Lipitor, Diazepam, Resveratrol.",
                   interactive=True)
  
      
  clear = gr.ClearButton([msg, chatbot])
  

  chat_btn.click(chat, [msg, tools], [msg, chatbot])
  gen_mols.select(add_gen, inputs = [gen_mols,percent_mask], outputs = msg)
  mols.select(add_mol, inputs = mols, outputs = msg)
  msg.submit(chat, [msg, tools], [msg, chatbot])
  clear.click(clear_history)

if __name__ == "__main__":
    forest.launch(debug=False, share=True)