import os
import json
import google.generativeai as genai
import gradio as gr

# Configure the Gemini API
genai.configure(api_key=os.environ["GEMINI_API_KEY"])

# Define the system instruction (pre_prompt)
pre_prompt = """
You are a chemistry expert. Break down the given chemical reaction mechanism into simple steps. 
The output should be in JSON format with the following structure:

{
    "step 1": {
        "reactants": ["reactant 1", "reactant 2"],
        "products": ["product 1", "product 2"],
        "mechanism": "Describe the perfectly logical reason behind this step, such as driving force, why bonds breaking, why bonds forming, electron transfer, what caused it.",
        "reagent": "Optional reagent or conditions for this step",
        "conditions": "Optional environmental conditions like temperature and pressure for this step"
    },
    "step 2": { ... }
    ...
}

DO NOT USE CODEBLOCK or any markdown. Simply write the JSON only.
"""

# Define the model with the system instruction (pre_prompt)
model = genai.GenerativeModel(
    model_name="gemini-1.5-flash",
    system_instruction=pre_prompt
)

# Function to generate steps for A -> D flow
def generate_reaction_steps(reactants, products):
    prompt = f"Given reactants: {reactants} and products: {products}, break down the reaction mechanism into simple steps in JSON format."
    
    chat_session = model.start_chat(history=[])
    response = chat_session.send_message(prompt)
    
    # Extract the JSON content from the response
    try:
        # Parsing the raw response to extract the JSON text
        content = response.text
        print(response)
        print("\n\n\n")
        print(content)
        
        # Loading the JSON string to a Python dictionary
        steps = json.loads(content)
    except (json.JSONDecodeError, KeyError):
        steps = {"error": "Failed to decode JSON from Gemini response."}
    
    return steps


# Gradio interface
def process_reaction(reactants, products):
    steps = generate_reaction_steps(reactants, products)
    return json.dumps(steps, indent=4)

# Create the Gradio interface
iface = gr.Interface(
    fn=process_reaction,
    inputs=[gr.Textbox(label="Reactants (comma-separated)"), gr.Textbox(label="Products (comma-separated)")],
    outputs="json",
    title="Chemistry Rationalizer",
    description="Break down a reaction mechanism into simple steps."
)

if __name__ == "__main__":
    iface.launch()