app.py

import streamlit as st
from azure.cosmos import CosmosClient, exceptions
import os
import pandas as pd
import json
from datetime import datetime
import shutil
from github import Github
from git import Repo
import base64

# Cosmos DB configuration
ENDPOINT = "https://acae-afd.documents.azure.com:443/"
SUBSCRIPTION_ID = "003fba60-5b3f-48f4-ab36-3ed11bc40816"
DATABASE_NAME = os.environ.get("COSMOS_DATABASE_NAME")
Key = os.environ.get("Key")

# GitHub configuration
def download_github_repo(url, local_path):
    if os.path.exists(local_path):
        shutil.rmtree(local_path)
    Repo.clone_from(url, local_path)

def create_zip_file(source_dir, output_filename):
    shutil.make_archive(output_filename, 'zip', source_dir)

def create_repo(g, repo_name):
    user = g.get_user()
    return user.create_repo(repo_name)

def push_to_github(local_path, repo, github_token):
    repo_url = f"https://{github_token}@github.com/{repo.full_name}.git"
    local_repo = Repo(local_path)
    
    if 'origin' in [remote.name for remote in local_repo.remotes]:
        origin = local_repo.remote('origin')
        origin.set_url(repo_url)
    else:
        origin = local_repo.create_remote('origin', repo_url)
    
    if not local_repo.heads:
        local_repo.git.checkout('-b', 'main')
        current_branch = 'main'
    else:
        current_branch = local_repo.active_branch.name
    
    local_repo.git.add(A=True)
    
    if local_repo.is_dirty():
        local_repo.git.commit('-m', 'Initial commit')
    
    origin.push(refspec=f'{current_branch}:{current_branch}')

def get_base64_download_link(file_path, file_name):
    with open(file_path, "rb") as file:
        contents = file.read()
    base64_encoded = base64.b64encode(contents).decode()
    return f'<a href="data:application/zip;base64,{base64_encoded}" download="{file_name}">Download {file_name}</a>'

# Function to fetch all databases and containers
def fetch_db_structure(client):
    structure = {}
    for database in client.list_databases():
        db_name = database['id']
        structure[db_name] = []
        db_client = client.get_database_client(db_name)
        for container in db_client.list_containers():
            structure[db_name].append(container['id'])
    return structure

# Function to fetch items from a container
def fetch_items(client, db_name, container_name):
    container = client.get_database_client(db_name).get_container_client(container_name)
    items = list(container.read_all_items())
    return items

# Function to render sidebar
def render_sidebar(client, structure):
    st.sidebar.title("📁 Cosmos DB Manager")
    
    # Database and Container selection
    selected_db = st.sidebar.selectbox("🗄️ Select Database", list(structure.keys()))
    selected_container = st.sidebar.selectbox("📦 Select Container", structure[selected_db])
    
    # Fetch items for the selected container
    items = fetch_items(client, selected_db, selected_container)
    
    # Item selection
    item_ids = [item['id'] for item in items]
    selected_item = st.sidebar.selectbox("📄 Select Item", ["Create New Item"] + item_ids)
    
    # CRUD buttons
    if selected_item != "Create New Item":
        if st.sidebar.button("🖊️ Edit Item"):
            st.session_state.action = "edit"
            st.session_state.selected_item = selected_item
        if st.sidebar.button("🗑️ Delete Item"):
            st.session_state.action = "delete"
            st.session_state.selected_item = selected_item
    else:
        if st.sidebar.button("➕ Create New Item"):
            st.session_state.action = "create"
    
    # GitHub section
    st.sidebar.subheader("🐙 GitHub Operations")
    github_token = os.environ.get("GITHUB")
    source_repo = st.sidebar.text_input("🔗 Source Repo URL", value="https://github.com/AaronCWacker/AIExamples-8-24-Streamlit")
    new_repo_name = st.sidebar.text_input("📝 New Repo Name", value=f"AIExample-Clone-{datetime.now().strftime('%Y%m%d_%H%M%S')}")
    
    if st.sidebar.button("📥 Clone Repository"):
        clone_github_repo(github_token, source_repo, new_repo_name)
    
    if st.sidebar.button("📤 Push to New Repository"):
        push_to_new_repo(github_token, source_repo, new_repo_name)
    
    # Archive data button
    if st.sidebar.button("📦 Archive All Data"):
        download_link = archive_all_data(client)
        st.sidebar.markdown(download_link, unsafe_allow_html=True)
    
    # Logout button
    if st.sidebar.button("🚪 Logout"):
        st.session_state.logged_in = False
        st.session_state.action = None
        st.session_state.selected_item = None
        st.rerun()
    
    return selected_db, selected_container, selected_item

# GitHub operations
def clone_github_repo(github_token, source_repo, new_repo_name):
    if github_token and source_repo:
        try:
            local_path = f"./temp_repo_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            download_github_repo(source_repo, local_path)
            zip_filename = f"{new_repo_name}.zip"
            create_zip_file(local_path, zip_filename[:-4])
            st.sidebar.markdown(get_base64_download_link(zip_filename, zip_filename), unsafe_allow_html=True)
            st.sidebar.success("Repository cloned successfully!")
        except Exception as e:
            st.sidebar.error(f"An error occurred: {str(e)}")
        finally:
            if os.path.exists(local_path):
                shutil.rmtree(local_path)
            if os.path.exists(zip_filename):
                os.remove(zip_filename)
    else:
        st.sidebar.error("Please ensure GitHub token is set in environment variables and source repository URL is provided.")

def push_to_new_repo(github_token, source_repo, new_repo_name):
    if github_token and source_repo:
        try:
            g = Github(github_token)
            new_repo = create_repo(g, new_repo_name)
            local_path = f"./temp_repo_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            download_github_repo(source_repo, local_path)
            push_to_github(local_path, new_repo, github_token)
            st.sidebar.success(f"Repository pushed successfully to {new_repo.html_url}")
        except Exception as e:
            st.sidebar.error(f"An error occurred: {str(e)}")
        finally:
            if os.path.exists(local_path):
                shutil.rmtree(local_path)
    else:
        st.sidebar.error("Please ensure GitHub token is set in environment variables and source repository URL is provided.")

# Function to archive all data
def archive_all_data(client):
    try:
        base_dir = "./cosmos_archive"
        if os.path.exists(base_dir):
            shutil.rmtree(base_dir)
        os.makedirs(base_dir)

        for database in client.list_databases():
            db_name = database['id']
            db_dir = os.path.join(base_dir, db_name)
            os.makedirs(db_dir)

            db_client = client.get_database_client(db_name)
            for container in db_client.list_containers():
                container_name = container['id']
                container_dir = os.path.join(db_dir, container_name)
                os.makedirs(container_dir)

                container_client = db_client.get_container_client(container_name)
                items = list(container_client.read_all_items())

                with open(os.path.join(container_dir, f"{container_name}.json"), 'w') as f:
                    json.dump(items, f, indent=2)

        archive_name = f"cosmos_archive_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        shutil.make_archive(archive_name, 'zip', base_dir)
        
        return get_base64_download_link(f"{archive_name}.zip", f"{archive_name}.zip")
    except Exception as e:
        return f"An error occurred while archiving data: {str(e)}"

# Main Streamlit app
def main():
    st.set_page_config(layout="wide")
    st.title("🌟 Cosmos DB and GitHub Integration")

    # Initialize session state
    if 'logged_in' not in st.session_state:
        st.session_state.logged_in = False
    if 'action' not in st.session_state:
        st.session_state.action = None
    if 'selected_item' not in st.session_state:
        st.session_state.selected_item = None

    # Login section
    if not st.session_state.logged_in:
        st.subheader("🔐 Login")
        input_key = st.text_input("Enter your Cosmos DB key", type="password")
        if st.button("🚀 Login"):
            if input_key:
                st.session_state.primary_key = input_key
                st.session_state.logged_in = True
                st.rerun()
            else:
                st.error("Invalid key. Please enter a valid Cosmos DB key.")
    else:
        # Initialize Cosmos DB client
        try:
            client = CosmosClient(ENDPOINT, credential=st.session_state.primary_key)
            
            # Fetch database structure
            structure = fetch_db_structure(client)
            
            # Render sidebar and get selections
            selected_db, selected_container, selected_item = render_sidebar(client, structure)
            
            # Main area
            if st.session_state.action == "create":
                create_item(client, selected_db, selected_container)
            elif st.session_state.action == "edit":
                edit_item(client, selected_db, selected_container, st.session_state.selected_item)
            elif st.session_state.action == "delete":
                delete_item(client, selected_db, selected_container, st.session_state.selected_item)
            else:
                st.write("Select an action from the sidebar to get started.")

        except exceptions.CosmosHttpResponseError as e:
            st.error(f"Failed to connect to Cosmos DB. HTTP error: {str(e)}. Status code: {e.status_code}")
        except Exception as e:
            st.error(f"An unexpected error occurred: {str(e)}")

if __name__ == "__main__":
    main()