import streamlit as st
import json
from azure.cosmos import CosmosClient

# Initialize variables
connection_str = ""
database_name = ""
container_name = ""

# 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)

# Connect to Azure Cosmos DB
def connect_to_cosmosdb(url, key, database, container):
    client = CosmosClient(url, credential=key)
    db = client.get_database_client(database)
    container = db.get_container_client(container)
    return container

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

# Connection Details Expander
with st.expander("Connect 🌍"):
    connection_str = st.text_input("Connection String (URL):", connection_str)
    database_name = st.text_input("Database Name:", database_name)
    container_name = st.text_input("Container Name:", container_name)
    if st.button("Connect"):
        try:
            container = connect_to_cosmosdb(connection_str, database_name, 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 🚀"):
        try:
            results = list(container.query_items(query=query, enable_cross_partition_query=True))
            st.write("Results 📋:")
            st.json(results)
        except Exception as e:
            st.error(f"Query failed: {e}")