import os
import streamlit as st
from azure.cosmos import CosmosClient, PartitionKey
from azure.cosmos.exceptions import CosmosResourceNotFoundError

# Initialize Azure Cosmos DB Client
COSMOS_CONNECTION_STRING = os.getenv('COSMOS_CONNECTION_STRING')
cosmos_client = CosmosClient.from_connection_string(COSMOS_CONNECTION_STRING)

# Function to Delete an Item in Cosmos DB
def delete_cosmos_item(db_name, container_name, item_id, partition_key):
    try:
        database_client = cosmos_client.get_database_client(db_name)
        container_client = database_client.get_container_client(container_name)
        container_client.delete_item(item=item_id, partition_key='id')
        st.success(f"Deleted Item: {item_id}")
    except CosmosResourceNotFoundError:
        st.error(f"Item not found: {item_id}")

# Display and Manage Cosmos DB Structure
def display_and_manage_cosmos_db():
    st.subheader('Azure Cosmos DB Structure')
    for db_properties in cosmos_client.list_databases():
        db_name = db_properties['id']
        database_client = cosmos_client.get_database_client(db_name)
        for container_properties in database_client.list_containers():
            container_name = container_properties['id']
            container_client = database_client.get_container_client(container_name)
            for item in container_client.read_all_items():
                st.json(item)  # Display the full data of the item
                if 'file_name' in item and os.path.isfile(item['file_name']):
                    st.image(item['file_name'])  # Show image if available
                # Assume the partition key is the item ID
                partition_key = item.get('id')
                delete_button_key = f"delete_{item['id']}"
                if st.button(f"🗑️ Delete {item['id']}", key=delete_button_key):
                    delete_cosmos_item(db_name, container_name, item['id'], partition_key)

# Main UI Function
def main():
    display_and_manage_cosmos_db()

# Run the main function
if __name__ == "__main__":
    main()