from pathlib import Path
import streamlit as st
import streamlit.components.v1 as components
from PIL import Image
import base64
import textgrad as tg
from stqdm import stqdm
import json
import os

def read_markdown_file(markdown_file):
    return Path(markdown_file).read_text()

def render_svg(svg_filename):
    with open(svg_filename,"r") as f:
        lines = f.readlines()
        svg=''.join(lines)
    """Renders the given svg string."""
    b64 = base64.b64encode(svg.encode('utf-8')).decode("utf-8")
    html = r'<img src="data:image/svg+xml;base64,%s"/>' % b64
    st.write(html, unsafe_allow_html=True)

def load_json(file_path):
    with open(file_path, "r") as f:
        return json.load(f)

def app():
    st.title('TextGrad: Automatic ''Differentiation'' via Text')
    st.image('assets/img/textgrad_logo_1x.png', width=200)

    st.markdown("### Examples")
    examples = {
        "Code Editor": load_json("examples/code_editor.json"),
        "Example Math Solution": load_json("examples/example_math.json"),
        # "🏞️ Image QA": load_json("examples/mathvista.json"),
    }
    
    col1, col2 = st.columns([1,1])
    with col1:
        # Display clickable examples
        option = st.selectbox(
            "Select an example to use:",
            examples.keys()
            )

        # Initialize session state if not already done
        st.session_state.data = examples[option]
        # if st.button("Clear example selection"):
        #     st.session_state.data = ""
        #     option = "None"

        # Initialize session state if not already done
        if "selected_option" not in st.session_state:
            st.session_state.selected_option = None
            st.session_state.object = None

        # Check if the selected option has changed
        if st.session_state.selected_option != option:
            st.session_state.data = examples[option]
            st.session_state.selected_option = option
            st.session_state.object_initialized = False

        # Display the selected option
        st.markdown(f"**Example selected:** {option}")
        
    with col2:
        valid_api_key = st.session_state.get('valid_api_key', False)
        form = st.form(key='api_key')
        OPENAI_API_KEY = form.text_input(label='Enter OpenAI API Key. Keys will not be stored', type = 'password')
        submit_button = form.form_submit_button(label='Submit')

        if submit_button:
            os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
           
            try:
                engine = tg.get_engine("gpt-4o")
                response = engine.generate("hello, what model is this?")
                st.session_state.valid_api_key = True
                valid_api_key = True
            except:
                st.error("Please enter a valid OpenAI API key.")
                st.stop()

    
    if valid_api_key:
        # Initialize the object when the option is selected
        if not st.session_state.object_initialized:
            if option == "Example Math Solution":
                from examples import example_math_scripts
                st.session_state.iteration = 0
                st.session_state.object = example_math_scripts.MathSolution(data=st.session_state.get("data", ""))
            elif option == "Code Editor":
                from examples import code_editor_scripts
                st.session_state.iteration = 0
                st.session_state.object = code_editor_scripts.CodeEditor(data=st.session_state.get("data", ""))
            # elif option == "🏞️ Image QA":
            #     from examples import example_imageQA_scripts
            #     st.session_state.iteration = 0
            #     st.session_state.object = example_imageQA_scripts.ImageQA(data=st.session_state.get("data", ""))
            st.session_state.object_initialized = True

        # Ensure object is loaded
        st.session_state.object.load_layout()

        # Run TextGrad
        # Add custom CSS to style the button
        # st.markdown("""
        #     <style>
        #     .large-button {
        #         font-size: 24px;
        #         padding: 16px 32px;
        #         background-color: #007bff;
        #         color: white;
        #         border: none;
        #         border-radius: 4px;
        #         cursor: pointer;
        #     }
        #     .large-button:hover {
        #         background-color: #0056b3;
        #     }
        #     </style>
        #     """, unsafe_allow_html=True)

        # # Create a larger button using HTML
        # if st.markdown('<button class="large-button">Run 1 iteration of TextGrad</button>', unsafe_allow_html=True):
        if st.button("Run 1 iteration of TextGrad"):
            st.session_state.object.show_results()


    st.markdown("""---""")
    st.markdown('### Disclaimer')
    st.markdown("Do not get addicted to TextGrad. It is a powerful tool that can be used for good or evil. Use it responsibly.")