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

# --- Load Config ---
def load_config(config_file="catalog_db_filepaths.yaml"):
    with open(config_file, 'r') as file:
        config = yaml.safe_load(file)
    return config

# --- Check if DuckDB Exists ---
def check_db_exists(db_filepath):
    return os.path.exists(db_filepath)

# --- List Available Tables ---
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 []

# --- Get Available Days ---
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 []

def get_schema_from_db(db_filepath, selected_table):
    """
    Fetches the schema from the DuckDB database for the specified table.

    Parameters:
    - db_filepath: Path to the DuckDB database.
    - selected_table: Name of the table to fetch the schema from.

    Returns:
    - List of field names in the table schema.
    """
    if check_db_exists(db_filepath):
        con = duckdb.connect(db_filepath)
        query = f"DESCRIBE {selected_table};"
        schema = con.execute(query).fetchdf()
        con.close()
        return schema['column_name'].tolist()  # Assuming DuckDB returns 'column_name'
    return []

def get_roi_fields_from_schema(schema):
    """
    Extracts ROI-related fields from the schema.
    
    Parameters:
    - schema: List of field names from the table schema.

    Returns:
    - List of field names containing '{ROI}_has_snow_presence' and '{ROI}_iflag_'.
    """
    # Extract fields that match the patterns
    roi_fields = [field for field in schema if 'has_snow_presence' in field or 'iflag_' in field]
    return roi_fields



# --- Load Image Data ---
def load_data(db_filepath, selected_table, selected_year, day_of_year, is_L1):
    """
    Loads data from the DuckDB database, including ROI fields dynamically based on schema.

    Parameters:
    - db_filepath: Path to the DuckDB database.
    - selected_table: Name of the table to query.
    - selected_year: Year to filter by.
    - day_of_year: Day of the year to filter by.
    - is_L1: Boolean to filter on is_L1.

    Returns:
    - DataFrame containing the query result.
    """
    # Get schema from the database
    schema = get_schema_from_db(db_filepath, selected_table)
    
    # Extract ROI-related fields
    roi_fields = get_roi_fields_from_schema(schema)
    
    # Base fields to always include in the query
    base_fields = [
        "catalog_guid", "year", "day_of_year", "catalog_filepath", "creation_date", "is_ready_for_products_use"
    ]
    
    # Combine base fields with ROI fields
    selected_fields = base_fields + roi_fields
    
    # Build the dynamic SELECT query
    query = f"""
        SELECT {', '.join(selected_fields)}
        FROM {selected_table}
        WHERE year = {selected_year} AND day_of_year = '{day_of_year}' AND is_L1 = {is_L1}
        AND catalog_filepath IS NOT NULL
        ORDER BY catalog_filepath ASC;
    """
    
    # Execute the query and fetch the result as DataFrame
    if check_db_exists(db_filepath):
        con = duckdb.connect(db_filepath)
        data = con.execute(query).fetchdf()
        con.close()
        return data
    return pd.DataFrame()  # Return an empty DataFrame if database does not exist

# --- Update is_ready_for_products_use ---
def update_is_ready_for_products_use(db_filepath, selected_table, catalog_guid, new_value):
    query = f"""
        UPDATE {selected_table}
        SET is_ready_for_products_use = {new_value}
        WHERE catalog_guid = '{catalog_guid}';
    """
    if check_db_exists(db_filepath):
        try:
            con = duckdb.connect(db_filepath)
            con.execute(query)
            con.close()
            return True  # Update successful
        except Exception as e:
            st.error(f"Failed to update the database: {e}")
            return False  # Update failed
    return False  # Database check failed

# --- Set All to Ready for Product Use ---
def set_all_ready_for_products_use(db_filepath, selected_table, image_data):
    for idx in range(len(image_data)):
        catalog_guid = image_data.iloc[idx]['catalog_guid']
        update_is_ready_for_products_use(db_filepath, selected_table, catalog_guid, True)
    #st.toast("All images marked as ready for product use.")

# --- Reload Image Data ---
def reload_image_data(catalog_db_filepath, selected_table, selected_year, day_of_year, is_L1):
    image_data = load_data(catalog_db_filepath, selected_table, selected_year, day_of_year, is_L1)
    st.session_state['image_data'] = image_data
    return image_data

# --- Render Image Grid ---
def render_image_grid(image_data, db_filepath, selected_table):
    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']
                is_ready_for_products_use = image_data.iloc[image_idx]['is_ready_for_products_use']

                # Display the image and checkbox for is_ready_for_products_use
                with col:
                    if os.path.exists(image_path):
                        img = Image.open(image_path)
                        st.image(img, use_column_width=True)
                        st.caption(f"Creation Date: {creation_date}")
                        checkbox_value = st.checkbox(f"Ready for Product Use", value=is_ready_for_products_use, key=catalog_guid)
                        if checkbox_value != is_ready_for_products_use:
                            update_is_ready_for_products_use(db_filepath, selected_table, catalog_guid, checkbox_value)

# --- Render Buttons and Update day_of_year ---
def render_buttons_and_update_day(available_days, day_of_year, db_filepath, selected_table, image_data):
    """
    Renders buttons to navigate days and set snow presence for all images.
    
    Parameters:
    - available_days: List of available days for the selected year.
    - day_of_year: The currently selected day of the year.
    - db_filepath: Path to the DuckDB database.
    - selected_table: The selected table name in the database.
    - image_data: The DataFrame containing image metadata.
    """
    _, col1, col2, col3 = st.columns([5, 1, 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]
            st.session_state['current']['day_of_year'] = day_of_year
            st.session_state['image_data'] = reload_image_data(db_filepath, selected_table, st.session_state['selected_year'], day_of_year, True)
    
    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]
            st.session_state['current']['day_of_year'] = day_of_year
            st.session_state['image_data'] = reload_image_data(db_filepath, selected_table, st.session_state['selected_year'], day_of_year, True)
    
    with col3:
        # Set all 'has_snow_presence' to True
        #if st.button("Set All Snow Presence to True"):
        #    set_all_snow_presence_to_true(db_filepath, selected_table, image_data)
        #    st.session_state['image_data'] = reload_image_data(db_filepath, selected_table, st.session_state['selected_year'], day_of_year, True)

        # Set all 'is_ready_for_products_use' to True
        if st.button("Set All to Ready"):
            set_all_ready_for_products_use(db_filepath, selected_table, image_data)
            st.session_state['image_data'] = reload_image_data(db_filepath, selected_table, st.session_state['selected_year'], day_of_year, True)


# ROIS 

def extract_roi_fields(data):
    """
    Extracts ROI-related fields from the loaded data and organizes them into a dictionary.
    
    Parameters:
    - data: DataFrame containing image metadata.

    Returns:
    - Dictionary structure: 
      {catalog_guid: {ROI: {'has_snow_presence': ..., 'iflag_fields': {...}}}}
    """
    result = {}

    # Filter the columns that match {ROI_01}_has_snow_presence and {ROI_01}_iflag_ patterns
    roi_columns = [col for col in data.columns if 'has_snow_presence' in col or 'iflag_' in col]

    # Iterate over each row in the data (each image)
    for _, row in data.iterrows():
        catalog_guid = row['catalog_guid']
        result[catalog_guid] = {}

        # Identify all unique ROIs by looking at the column prefixes
        rois = set(col.split('_')[0] for col in roi_columns)

        # Iterate through each unique ROI
        for roi in rois:
            roi_data = {}
            
            # For each ROI, grab the related fields
            snow_presence_field = f"{roi}_has_snow_presence"
            iflag_fields = {col: row[col] for col in roi_columns if col.startswith(f"{roi}_iflag_")}

            # Check if snow presence field exists
            if snow_presence_field in data.columns:
                roi_data['has_snow_presence'] = row[snow_presence_field]

            # Add iflag fields to the ROI data
            roi_data['iflag_fields'] = iflag_fields

            # Add the ROI data to the result dictionary
            result[catalog_guid][roi] = roi_data

    return result


# --- Main Application Logic ---
def main():
    st.title("Phenocam Image Viewer")

    # Initialize session state for 'current' if not already done
    if 'current' not in st.session_state:
        st.session_state['current'] = {
            'station_name': None,
            'day_of_year': None
        }

    # Load the config
    config = load_config()

    # Header columns
    hcol1, hcol2 = st.columns([2, 2])
    
    # Select station
    with hcol1:
        station_names = list(config.keys())
        selected_station = st.selectbox("Select Station", station_names)
    
    # Get catalog DB filepath
    db_dirpath = config[selected_station]['db_dirpath']
    db_filename = config[selected_station]['db_filename']
    catalog_db_filepath = os.path.join(db_dirpath, db_filename)
    st.session_state['db_filepath'] = catalog_db_filepath
    st.session_state['current']['station_name'] = selected_station

    # Select table
    available_tables = list_tables(catalog_db_filepath)
    with hcol2:
        selected_table = st.selectbox("Select Table", available_tables)
    st.session_state['selected_table'] = selected_table
    
    # Select year
    with hcol1:
        selected_year = st.selectbox("Select Year", [2022, 2023, 2024])
    st.session_state['selected_year'] = selected_year

    # Get available days
    available_days = get_available_days(catalog_db_filepath, selected_table, selected_year, True)

    if available_days:
        # Initialize day_of_year if it's not set or is None
        if st.session_state['current']['day_of_year'] is None:
            st.session_state['current']['day_of_year'] = available_days[0]  # Set to the first available day
        
        # Load the image data
        day_of_year = st.session_state['current']['day_of_year']
        image_data = load_data(catalog_db_filepath, selected_table, selected_year, day_of_year, True)
        
        # Render buttons for previous/next day and set all to ready
        render_buttons_and_update_day(available_days, day_of_year, catalog_db_filepath, selected_table, image_data)

        # Sync the selectbox with the updated day_of_year
        with hcol2:
            # Ensure day_of_year is valid in available_days
            if day_of_year not in available_days:
                day_of_year = available_days[0]
                st.session_state['current']['day_of_year'] = day_of_year
            
            day_of_year = st.selectbox("Day of Year", available_days, index=available_days.index(st.session_state['current']['day_of_year']))
            if day_of_year != st.session_state['current']['day_of_year']:
                st.session_state['current']['day_of_year'] = day_of_year
                st.session_state['image_data'] = reload_image_data(catalog_db_filepath, selected_table, selected_year, day_of_year, True)

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

        # Extract ROI-related fields from the loaded data
        #roi_fields = extract_roi_fields(image_data)
        st.data_editor(image_data)
    else:
        st.write(f"No images found for year {selected_year}.")

if __name__ == "__main__":
   main()
