app.py

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 Cosmos DB Client
cosmos_client = CosmosClient.from_connection_string(COSMOS_CONNECTION_STRING)

# Initialize Azure Blob Storage Client
blob_service = BlobServiceClient.from_connection_string(BLOB_STORAGE_CONNECTION_STRING)


# Function to Retrieve and Display Cosmos DB Structure
def display_cosmos_db_structure():
    st.subheader('Azure Cosmos DB Structure')
    for db_properties in cosmos_client.list_databases():
        db_name = db_properties['id']
        st.markdown(f"#### Database: {db_name}")
        database_client = cosmos_client.get_database_client(db_name)
        
        for container_properties in database_client.list_containers():
            container_name = container_properties['id']
            st.markdown(f"- **Container**: {container_name}")
            container_client = database_client.get_container_client(container_name)
            
            items = list(container_client.read_all_items())
            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)

# Function to Add or Update an Item
def add_or_update_item(database_name, container_name, item_id, item_data):
    container = cosmos_client.get_database_client(database_name).get_container_client(container_name)
    query = "SELECT * FROM c WHERE c.id = @id"
    parameters = [{"name": "@id", "value": item_id}]
    items = list(container.query_items(query=query, parameters=parameters, enable_cross_partition_query=True))

    if items:
        # Item exists, update it
        existing_item = items[0]
        existing_item.update(item_data)
        container.replace_item(item=existing_item, body=existing_item)
        st.write(f"Updated Item: {item_id}")
    else:
        # Item does not exist, create new
        container.create_item(body=item_data)
        st.write(f"Created New Item: {item_id}")

    # Display item ID and image if available
    display_item(database_name, container_name, item_id)

# Function to Display an Item
def display_item(database_name, container_name, item_id):
    container = cosmos_client.get_database_client(database_name).get_container_client(container_name)
    item = container.read_item(item_id, partition_key=PartitionKey(item_id))
    st.markdown(f"- **Item ID**: {item_id}")
    if 'file_name' in item and item['file_name'].endswith('.png'):
        st.image(item['file_name'], caption=f"Image: {item['file_name']}")

                    
# Button to Trigger Display of Cosmos DB Structure
if st.button('Show Cosmos DB Structure'):
    display_cosmos_db_structure()
        
# Test Function to Insert PNG Images
def test_insert_png_images():
    # Get the first database and container
    db_properties = next(cosmos_client.list_databases(), None)
    if db_properties:
        db_name = db_properties['id']
        database_client = cosmos_client.get_database_client(db_name)
        container_properties = next(database_client.list_containers(), None)
        if container_properties:
            container_name = container_properties['id']

            # Insert PNG files
            png_files = glob.glob('*.png')
            for file_name in png_files:
                item_id = os.path.splitext(file_name)[0]  # Use file name without extension as ID
                item_data = {"id": item_id, "file_name": file_name}
                add_or_update_item(database_client, container_name, item_id, item_data)
                st.write(f"Inserted: {file_name}")

            # Displaying Images
            st.subheader("Displaying Images in Container")
            items = list(database_client.get_container_client(container_name).read_all_items())
            for item in items:
                if 'file_name' in item and item['file_name'].endswith('.png'):
                    st.image(item['file_name'], caption=item['file_name'])
        else:
            st.error("No container found in the database.")
    else:
        st.error("No database found.")

# UI to Test Image Insertion Function
if st.button('Test Insert PNG Images'):
    test_insert_png_images()



# Function to Add an Item
def add_item_to_container(database_name, container_name, item):
    container = cosmos_client.get_database_client(database_name).get_container_client(container_name)
    container.create_item(item)

# Function to Get All Item Keys from a Container
def get_all_item_keys(database_name, container_name):
    container = cosmos_client.get_database_client(database_name).get_container_client(container_name)
    items = list(container.query_items(
        query="SELECT c.id FROM c",
        enable_cross_partition_query=True))
    return [item['id'] for item in items]

# Function to Retrieve and Display an Item
def display_item(database_name, container_name, item_id):
    container = cosmos_client.get_database_client(database_name).get_container_client(container_name)
    item = container.read_item(item_id, partition_key=PartitionKey(item_id))
    st.markdown(f"- **Item ID**: {item_id}, **Content**: {json.dumps(item)}")

# Test Function to Add and Retrieve Image Prompts
def test_image_prompts(database_name, container_name):
    # Sample data
    image_prompts = [{"id": f"image_{i}", "prompt": f"Image prompt {i}", "image_url": f"http://example.com/image_{i}.jpg"} for i in range(10)]
    
    # Add items
    for item in image_prompts:
        add_item_to_container(database_name, container_name, item)

    # Retrieve and display items
    item_keys = get_all_item_keys(database_name, container_name)
    for key in item_keys:
        display_item(database_name, container_name, key)

# UI to Test Image Prompts Function
if st.button('Test Image Prompts'):
    test_database_name = st.text_input('Enter Test Database Name')
    test_container_name = st.text_input('Enter Test Container Name')
    if test_database_name and test_container_name:
        test_image_prompts(test_database_name, test_container_name)
    else:
        st.error('Please enter the test database and container names.')



    
# Streamlit UI
st.title('Azure Services Integration with Streamlit')

# Azure Cosmos DB - CRUD Operations
st.subheader('Azure Cosmos DB - CRUD Operations')
cosmos_db = st.text_input('Database Name')
cosmos_container = st.text_input('Container Name')
item_id = st.text_input("Item ID (for Read, Update, Delete)")
item_data = st.text_area("Item Data (JSON format, for Create and Update)")

if st.button('Create Item in Cosmos DB'):
    container = cosmos_client.get_database_client(cosmos_db).get_container_client(cosmos_container)
    container.create_item(item_data)

if st.button('Read Item from Cosmos DB'):
    container = cosmos_client.get_database_client(cosmos_db).get_container_client(cosmos_container)
    item = container.read_item(item_id, partition_key=PartitionKey(item_id))
    st.json(item)

# Azure Blob Storage - Upload/Download
st.subheader('Azure Blob Storage - Upload/Download')
blob_container = st.text_input('Blob Container')
blob_file = st.file_uploader('Upload file to Blob')

if blob_file is not None and st.button('Upload to Blob'):
    blob_client = blob_service.get_blob_client(container=blob_container, blob=blob_file.name)
    blob_client.upload_blob(blob_file.getvalue())

# Azure Functions - Trigger
st.subheader('Azure Functions - Trigger')
function_url = st.text_input('Function URL')

if st.button('Call Azure Function'):
    response = requests.get(function_url)
    st.write('Function Response:', response.text)