import os
import streamlit as st
from azure.cosmos import CosmosClient, PartitionKey
from azure.storage.blob import BlobServiceClient
from azure.cosmos.exceptions import CosmosResourceNotFoundError
import requests
import glob

# Environment Variables
COSMOS_CONNECTION_STRING = os.getenv('COSMOS_CONNECTION_STRING')
BLOB_STORAGE_CONNECTION_STRING = os.getenv('BLOB_STORAGE_CONNECTION_STRING')

# Initialize Azure Clients
cosmos_client = CosmosClient.from_connection_string(COSMOS_CONNECTION_STRING)
blob_service = BlobServiceClient.from_connection_string(BLOB_STORAGE_CONNECTION_STRING)

@st.cache(allow_output_mutation=True)
def get_cosmos_client():
    return cosmos_client

def get_database_client(db_name):
    cosmos_client = get_cosmos_client()
    return cosmos_client.get_database_client(db_name)

def get_container_client(db_name, container_name):
    database_client = get_database_client(db_name)
    return database_client.get_container_client(container_name)

def list_databases():
    cosmos_client = get_cosmos_client()
    return list(cosmos_client.list_databases())

def list_containers(db_name):
    database_client = get_database_client(db_name)
    return list(database_client.list_containers())

def list_items(db_name, container_name):
    container_client = get_container_client(db_name, container_name)
    return list(container_client.read_all_items())

def create_item(db_name, container_name, item_data):
    container_client = get_container_client(db_name, container_name)
    container_client.create_item(body=item_data)

def update_item(db_name, container_name, item_id, new_data):
    container_client = get_container_client(db_name, container_name)
    try:
        item = container_client.read_item(item_id, partition_key=item_id)
        for key in new_data:
            item[key] = new_data[key]
        container_client.replace_item(item_id, item)
    except CosmosResourceNotFoundError:
        st.error("Item not found for update")

def delete_item(db_name, container_name, item_id):
    container_client = get_container_client(db_name, container_name)
    try:
        container_client.delete_item(item_id, partition_key=item_id)
    except CosmosResourceNotFoundError:
        st.error("Item not found for deletion")

def display_cosmos_db_structure():
    st.subheader('Azure Cosmos DB Structure')

    # Fetching databases
    databases = list_databases()
    if not databases:
        st.error("No databases found.")
        return

    # Iterating through databases
    for db_properties in databases:
        db_name = db_properties['id']
        st.markdown(f"#### Database: {db_name}")

        # Fetching containers in the database
        containers = list_containers(db_name)
        if not containers:
            st.markdown("No containers found in this database.")
            continue

        # Iterating through containers
        for container_properties in containers:
            container_name = container_properties['id']
            st.markdown(f"- **Container**: {container_name}")

            # Fetching and displaying items in the container
            items = list_items(db_name, container_name)
            if not items:
                st.markdown("  - No items in this container.")
                continue

            for item in items:
                item_desc = f"  - Item: `{item['id']}`"
                if 'file_name' in item and item['file_name'].endswith('.png'):
                    st.markdown(item_desc)
                    st.image(item['file_name'])
                else:
                    st.markdown(item_desc)


def manage_png_images():
    # Get database and container names
    db_properties = next(cosmos_client.list_databases(), None)
    if not db_properties:
        st.error("No database found.")
        return

    db_name = db_properties['id']
    database_client = get_database_client(db_name)
    container_properties = next(database_client.list_containers(), None)
    if not container_properties:
        st.error("No container found.")
        return

    container_name = container_properties['id']
    container_client = get_container_client(db_name, container_name)

    # List existing items
    existing_items = list_items(db_name, container_name)
    existing_ids = {item['id']: item for item in existing_items}

    # Display existing items
    st.subheader("Existing PNG Images")
    for item_id, item in existing_ids.items():
        st.markdown(f"**{item_id}**: {item.get('file_name', 'No file name')}")
        if st.button(f"🗑️ Delete {item_id}", key=f"delete_{item_id}"):
            delete_item(db_name, container_name, item_id)
            st.success(f"Deleted Item: {item_id}")

    # Add or Update PNG files
    st.subheader("Add or Update PNG Images")
    png_files = glob.glob('*.png')
    for file_name in png_files:
        item_id = os.path.splitext(file_name)[0]
        item_data = {"id": item_id, "file_name": file_name}
        if item_id not in existing_ids:
            if st.button(f"➕ Add {item_id}", key=f"add_{item_id}"):
                create_item(db_name, container_name, item_data)
                st.success(f"Added Item: {item_id}")
        else:
            if st.button(f"✏️ Update {item_id}", key=f"update_{item_id}"):
                update_item(db_name, container_name, item_id, {"file_name": file_name})
                st.success(f"Updated Item: {item_id}")


def list_blobs_in_container(container_name):
    container_client = blob_service.get_container_client(container_name)
    return list(container_client.list_blobs())

# UI Elements
if st.button('Show Cosmos DB Structure'):
    display_cosmos_db_structure()

if st.button('Manage PNG Images'):
    manage_png_images()

# Listing Blobs in a Container
blob_container = st.text_input('Enter Blob Container Name')
if st.button('List Blobs in Container'):
    blobs = list_blobs_in_container(blob_container)
    for blob in blobs:
        st.write(blob.name)

# Remaining UI for Azure Blob Storage and Azure Functions as in the original code