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

# 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 Item by ID in Cosmos DB
def delete_item_by_id(db_name, container_name, item_id):
    database_client = cosmos_client.get_database_client(db_name)
    container_client = database_client.get_container_client(container_name)
    try:
        # Using the item ID as the partition key value
        container_client.delete_item(item=item_id, partition_key=item_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.markdown(f"- **Item ID**: `{item['id']}`")
                if st.button(f"🗑️ Delete {item['id']}", key=f"delete_{item['id']}"):
                    delete_item_by_id(db_name, container_name, item['id'])

# Insert PNG Images with Unique Identifiers
def insert_png_images_with_unique_ids():
    db_name, container_name = select_db_and_container()
    container_client = cosmos_client.get_database_client(db_name).get_container_client(container_name)
    for file_name in glob.glob('*.png'):
        unique_id = f"{os.path.splitext(file_name)[0]}_{datetime.now().strftime('%Y%m%d%H%M%S')}"
        item_data = {"id": unique_id, "file_name": file_name}
        container_client.create_item(body=item_data)
        st.write(f"Inserted Item: {unique_id}")

# Helper to Select Database and Container
def select_db_and_container():
    db_name = st.selectbox("Select Database", [db['id'] for db in cosmos_client.list_databases()])
    container_name = st.selectbox("Select Container", [container['id'] for container in cosmos_client.get_database_client(db_name).list_containers()])
    return db_name, container_name

# UI for Displaying and Managing Cosmos DB
display_and_manage_cosmos_db()

# UI for Inserting PNG Images
if st.button('Insert PNG Images with Unique IDs'):
    insert_png_images_with_unique_ids()