import streamlit as st
import duckdb
import pandas as pd
import yaml
import os
from PIL import Image

# Function to load config from yaml
def load_config(config_file="catalog_db_filepaths.yaml"):
    with open(config_file, 'r') as file:
        config = yaml.safe_load(file)
    return config

# Function to check if the DuckDB file exists
def check_db_exists(db_filepath):
    if not os.path.exists(db_filepath):
        st.error(f"Database file does not exist at: {db_filepath}")
        return False
    return True

# Function to get the list of available tables in the DuckDB database
def list_tables(db_filepath):
    if check_db_exists(db_filepath):
        con = duckdb.connect(db_filepath)
        tables = con.execute("SHOW TABLES;").fetchdf()
        con.close()
        return tables['name'].tolist()
    return []

# Function to load available days (day_of_year) where there are images
def get_available_days(db_filepath, selected_table, selected_year, is_L1):
    query = f"""
        SELECT DISTINCT day_of_year
        FROM {selected_table}
        WHERE year = {selected_year} AND is_L1 = {is_L1} AND catalog_filepath IS NOT NULL;
    """
    if check_db_exists(db_filepath):
        con = duckdb.connect(db_filepath)
        days = con.execute(query).fetchdf()
        con.close()
        return sorted(days['day_of_year'].tolist())  # Ensure the list is sorted
    return []

# Function to load data from DuckDB (filtered by year, day_of_year, and is_L1)
# Includes sorting of images by 'catalog_filepath' and fetching 'creation_date'
def load_data(db_filepath, selected_table, year, day_of_year, is_L1):
    query = f"""
        SELECT catalog_guid, year, day_of_year, catalog_filepath, creation_date, is_L1
        FROM {selected_table}
        WHERE year = {year} AND day_of_year = '{day_of_year}' AND is_L1 = {is_L1}
        AND catalog_filepath IS NOT NULL
        ORDER BY catalog_filepath ASC;  -- Ensure sorting by file path
    """
    if check_db_exists(db_filepath):
        con = duckdb.connect(db_filepath)
        data = con.execute(query).fetchdf()
        con.close()
        return data
    return pd.DataFrame()


# Class to manage image containers
class ImageContainer:
    def __init__(self, catalog_guid, image_path, creation_date):
        self.catalog_guid = catalog_guid
        self.image_path = image_path
        self.creation_date = creation_date

    # Function to display the image and its creation date if it exists
    def display(self):
        if os.path.exists(self.image_path):  # Ensure the image exists
            img = Image.open(self.image_path)
            st.image(img, use_column_width=True)
            st.caption(f"Creation Date: {self.creation_date}")
        else:
            st.write(f"Image not found: {self.image_path}")

# Function to render the grid of images dynamically based on available images
def render_image_grid(image_data):
    num_images = len(image_data)
    if num_images == 0:
        st.write("No images available.")
        return
    
    # Determine the number of columns (4 in this case) and rows based on images
    columns_per_row = 4
    num_rows = (num_images + columns_per_row - 1) // columns_per_row  # Round up to the nearest row
    
    for row_idx in range(num_rows):
        # Create the columns dynamically for the current row
        cols = st.columns(min(columns_per_row, num_images - row_idx * columns_per_row))
        
        for col_idx, col in enumerate(cols):
            image_idx = row_idx * columns_per_row + col_idx
            if image_idx < num_images:
                catalog_guid = image_data.iloc[image_idx]['catalog_guid']
                image_path = image_data.iloc[image_idx]['catalog_filepath']
                creation_date = image_data.iloc[image_idx]['creation_date']

                # Create an ImageContainer instance for each image
                img_container = ImageContainer(catalog_guid, image_path, creation_date)
                
                # Display the image in the respective column
                with col:
                    img_container.display()

# Initialize session state for the first run
if 'current' not in st.session_state:
    st.session_state['current'] = {
        'table': None,
        'year': 2024,  # Default year
        'day_of_year': 1,  # Default day of the year
        'station': None
    }

# Main Streamlit App
def main():
    st.title("Phenocam Image Viewer")

    # Load the config
    config = load_config()

    # header columns
    hcol1, hcol2 = st.columns(2)
    
    # Select the station
    station_names = list(config.keys())
    with hcol1:
        selected_station = st.selectbox("Select Station", station_names)
 
    
    # Get the catalog database filepath for the selected station
    db_dirpath = config[selected_station]['db_dirpath']
    db_filename = config[selected_station]['db_filename']
    catalog_db_filepath = os.path.join(db_dirpath, db_filename)

    # Store the selected station in session state
    st.session_state['current']['station_name'] = selected_station

    # List available tables and use selectbox to select a table
    available_tables = list_tables(catalog_db_filepath)
    with hcol2:
        selected_table = st.selectbox("Select Table", available_tables)
    
    is_L1 = True # st.checkbox("Is L1", value=True)
    
    # Store the selected table in session state
    st.session_state['current']['table_name'] = selected_table
    
     
    # User input to select the year and is_L1 flag
    with hcol1:
        selected_year = st.selectbox("Select Year", [2022, 2023, 2024])  # Update with actual year data
        

    # Get available days with images for the selected year
    available_days = get_available_days(catalog_db_filepath, selected_table, selected_year, is_L1)

    # Ensure the user can only select from available days
    if available_days:
        # If the current day is not in the available days, reset it to the first available day
        if st.session_state['current']['day_of_year'] not in available_days:
            st.session_state['current']['day_of_year'] = available_days[0]
        
        with hcol2:				
            # Display the day_of_year selectbox, synchronized with button clicks
            day_of_year = st.selectbox("Day of Year", available_days, index=available_days.index(st.session_state['current']['day_of_year']))
    else:
        st.write(f"No images found for year {selected_year} and is_L1={is_L1}.")
        return  # Stop execution if no images are available

    # Update session state with selected year and day_of_year
    st.session_state['current']['year'] = selected_year
    st.session_state['current']['day_of_year'] = day_of_year

    # Load the filtered data (sorted by catalog_filepath and includes creation_date)
    image_data = load_data(catalog_db_filepath, selected_table, selected_year, day_of_year, is_L1)

    # Display the grid of images using expander
    with st.expander(f"Images for Day {day_of_year} of {selected_year}", expanded=True):
        render_image_grid(image_data)

    # Add buttons for previous and next day navigation
    col1, col2 = st.columns([1, 1])
    
    with col1:
        if st.button("Previous Day") and day_of_year > min(available_days):
            day_of_year = available_days[available_days.index(day_of_year) - 1]
    
    with col2:
        if st.button("Next Day") and day_of_year < max(available_days):
            day_of_year = available_days[available_days.index(day_of_year) + 1]
    
    # Update session state with the new day_of_year
    st.session_state['current']['day_of_year'] = day_of_year

if __name__ == "__main__":
    main()
