import os
import gradio as gr
from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough, chain

def create_dynamic_chain(api_key):
    llm = ChatOpenAI(model="gpt-4o-mini", api_key=api_key)
    
    # Chain for general questions
    general_prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a helpful assistant that provides direct answers."),
        ("human", "{question}")
    ])
    
    # Chain for mathematical calculations
    math_prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a mathematical assistant. Solve the problem and show your work."),
        ("human", "{question}")
    ])
    
    # Chain for coding questions
    code_prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a coding assistant. Provide code examples and explanations."),
        ("human", "{question}")
    ])
    
    general_chain = general_prompt | llm | StrOutputParser()
    math_chain = math_prompt | llm | StrOutputParser()
    code_chain = code_prompt | llm | StrOutputParser()
    
    @chain
    def dynamic_chain(input_dict):
        question = input_dict["question"].lower()
        
        # Detect question type
        if any(word in question for word in ["calculate", "solve", "compute", "sum", "multiply"]):
            return math_chain
        elif any(word in question for word in ["code", "program", "function", "python", "javascript"]):
            return code_chain
        return general_chain

    return dynamic_chain

def process_message(message, history, api_key, example_select):
    if not api_key:
        return "", [{"role": "assistant", "content": "Please enter your OpenAI API key."}]
    
    try:
        # Handle example selection
        if example_select != "Custom Input":
            message = EXAMPLES[example_select]
        
        chain = create_dynamic_chain(api_key)
        response = chain.invoke({"question": message})
        
        history.append({"role": "user", "content": message})
        history.append({"role": "assistant", "content": response})
        
        return "", history
    except Exception as e:
        return "", history + [{"role": "assistant", "content": f"Error: {str(e)}"}]

# Example questions for different chain types
EXAMPLES = {
    "General Question": "What are the main features of renewable energy?",
    "Math Problem": "Calculate the area of a circle with radius 5 units.",
    "Coding Question": "Write a Python function to find the factorial of a number.",
    "Custom Input": ""
}

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Dynamic Chain Demo with Examples")
    
    with gr.Row():
        api_key = gr.Textbox(
            label="OpenAI API Key",
            placeholder="Enter your OpenAI API key",
            type="password"
        )
        example_select = gr.Dropdown(
            choices=list(EXAMPLES.keys()),
            value="Custom Input",
            label="Select Example"
        )
    
    chatbot = gr.Chatbot(type="messages")
    msg = gr.Textbox(label="Message", placeholder="Type your message or select an example above")
    clear = gr.ClearButton([msg, chatbot])
    
    # Example descriptions
    gr.Markdown("""
    ## Example Types:
    1. **General Questions**: Regular queries that don't require special processing
    2. **Math Problems**: Questions involving calculations and mathematical operations
    3. **Coding Questions**: Programming-related queries that return code examples
    
    ## Try these patterns:
    - Math: "Calculate...", "Solve...", "Compute..."
    - Code: "Write a function...", "Program...", "Code..."
    - General: Any other type of question
    """)
    
    msg.submit(
        process_message,
        inputs=[msg, chatbot, api_key, example_select],
        outputs=[msg, chatbot]
    )

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