import os
import pandas as pd
import streamlit as st
from azure.cosmos import CosmosClient
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):
    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}")

# Load Cosmos DB Data into a DataFrame
def load_cosmos_data_into_dataframe():
    data = []
    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():
                data.append(item)
    return pd.DataFrame(data)

# Display and Manage Cosmos DB Data
def display_and_manage_cosmos_data(df):
    st.subheader('Azure Cosmos DB Data')
    for index, row in df.iterrows():
        st.json(row.to_json())  # Display the full data of the item
        if 'file_name' in row and os.path.isfile(row['file_name']):
            st.image(row['file_name'])  # Show image if available
        delete_button_key = f"delete_{row['id']}"
        if st.button(f"🗑️ Delete {row['id']}", key=delete_button_key):
            delete_cosmos_item(row['_database_id'], row['_container_id'], row['id'])
            st.experimental_rerun()

# Main UI Function
def main():
    df = load_cosmos_data_into_dataframe()
    display_and_manage_cosmos_data(df)

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