import streamlit as st
import pandas as pd
import re
from azure.identity import DefaultAzureCredential
from azure.mgmt.compute import ComputeManagementClient

def clear_offers():
    if 'offers' in st.session_state: del st.session_state.offers
    if 'selected_offer' in st.session_state: del st.session_state.selected_offer
    clear_skus()

def clear_selected_sku():
    if 'selected_sku' in st.session_state: del st.session_state.selected_sku

def clear_skus():
    if 'skus' in st.session_state: del st.session_state.skus
    clear_selected_sku()

@st.cache_data
def get_publishers(location: str):
    return [publisher.name for publisher in compute_client.virtual_machine_images.list_publishers(location)]

@st.cache_data
def get_offers(location: str, publisher: str):
    return [offer.name for offer in compute_client.virtual_machine_images.list_offers(location, publisher)]

@st.cache_data
def get_skus(location: str, publisher: str, offer: str):
    return [sku.name for sku in compute_client.virtual_machine_images.list_skus(location, publisher, offer)]

@st.cache_data
def get_image_versions(location: str, publisher: str, offer: str, sku: str):
    return [version.name for version in compute_client.virtual_machine_images.list(location, publisher, offer, sku)]

subscription_id = "c885a276-c882-483f-b216-42f73715161d"
location = "westeurope"

credential = DefaultAzureCredential()
compute_client = ComputeManagementClient(credential, subscription_id)

st.set_page_config(page_title='Azure Image Chooser', layout='wide')

st.title('Azure Image Chooser')

left_col, middle_col, right_col = st.columns(3)

if 'publishers' not in st.session_state:
    st.session_state.publishers = get_publishers(location)
    clear_offers()

selected_publisher = left_col.selectbox('Select Publisher', options=st.session_state.publishers)

if 'selected_publisher' not in st.session_state or selected_publisher != st.session_state.selected_publisher:
    st.session_state.selected_publisher = selected_publisher
    clear_offers()

if 'offers' not in st.session_state:
    st.session_state.offers = get_offers(location, st.session_state.selected_publisher)
    clear_skus()

selected_offer = middle_col.selectbox('Select Offer', options=st.session_state.offers)

if 'selected_offer' not in st.session_state or selected_offer != st.session_state.selected_offer:
    st.session_state.selected_offer = selected_offer
    clear_skus()

if 'skus' not in st.session_state:
    st.session_state.skus = get_skus(location, st.session_state.selected_publisher, st.session_state.selected_offer)
    clear_selected_sku()

selected_sku = right_col.selectbox('Select SKU', options=st.session_state.skus)

if 'selected_sku' not in st.session_state or selected_sku != st.session_state.selected_sku:
    st.session_state.selected_sku = selected_sku

def version_key(v):
    return [int(x) for x in v.split('.')]

regex = re.compile(r'^[0-9]+\.[0-9]+\.[0-9]+$')

# Display available image versions
images_versions = get_image_versions(location, st.session_state.selected_publisher, st.session_state.selected_offer, st.session_state.selected_sku)

# Check if all image version string match the re.
if all(regex.match(version) for version in images_versions):
    images_versions = sorted(images_versions, key=version_key)

st.dataframe(images_versions, hide_index=True, column_config={"value": "Image Version"})