import streamlit as st
from azure.cosmos import CosmosClient

# Load and save query files
def load_query(filename):
    with open(filename, 'r') as file:
        return file.read()

def save_query(filename, query):
    with open(filename, 'w') as file:
        file.write(query)

# Streamlit UI
st.title("Azure Cosmos DB Explorer 👽")

# Default URI and KEY (You can set these values in code)
default_account_uri = "Your Default Cosmos DB URI Here"
default_account_key = "Your Default Cosmos DB Key Here"

client = None

# Connection Details Expander
with st.expander("Connect 🌍"):
    account_uri = st.text_input("Account URI:", default_account_uri)
    account_key = st.text_input("Account Key:", default_account_key, type="password")
    database_name = st.text_input("Database Name:", "")
    container_name = st.text_input("Container Name:", "")
    if st.button("Connect"):
        try:
            client = CosmosClient(account_uri, credential=account_key)
            database_client = client.get_database_client(database_name)
            container_client = database_client.get_container_client(container_name)
            st.success("Connected successfully! 🎉")
        except Exception as e:
            st.error(f"Failed to connect: {e}")

# Query Editor Expander
with st.expander("Query Editor 📝"):
    query = st.text_area("Enter your SQL query here:", "")
    file_option = st.selectbox("File Options", ["New", "Open", "Save", "Save As"])
    
    if file_option == "New":
        query = ""
    elif file_option == "Open":
        open_file = st.file_uploader("Choose a file:", type=["txt"])
        if open_file is not None:
            query = load_query(open_file)
    elif file_option == "Save":
        save_filename = st.text_input("Enter filename to save:", "my_query.txt")
        if st.button("Save Query"):
            save_query(save_filename, query)
    elif file_option == "Save As":
        saveas_filename = st.text_input("Enter new filename:", "my_new_query.txt")
        if st.button("Save As"):
            save_query(saveas_filename, query)
    
    if st.button("Execute Query 🚀"):
        if client:
            try:
                items = list(container_client.query_items(
                    query=query,
                    enable_cross_partition_query=True
                ))
                st.write("Results 📋:")
                st.json(items)
            except Exception as e:
                st.error(f"Query failed: {e}")
        else:
            st.warning("Not connected to any Cosmos DB. Please connect first.")

# Instructions
st.markdown("""
## Instructions to Run this App:

1. **Install Packages**: If you haven't, install the required Python packages:
2. **Run App**: Save this code in a file, say `streamlit_cosmosdb_app.py`, and then run `streamlit run streamlit_cosmosdb_app.py`.
3. **Execute**: Use the UI to connect and run SQL queries against your Cosmos DB.
""")