import streamlit as st
import boto3
from botocore.exceptions import NoCredentialsError, ClientError, EndpointConnectionError
from collections import defaultdict
from datetime import datetime

# AWS S3 setup
os.environ['AWS_ACCESS_KEY']
os.environ['AWS_SECRET_KEY']
os.environ['BUCKET_NAME']

# Initialize a session using DigitalOcean Spaces
session = boto3.session.Session()
s3 = session.client('s3',
                    region_name='us-west-2',  # Update to the correct region
                    aws_access_key_id=AWS_ACCESS_KEY,
                    aws_secret_access_key=AWS_SECRET_KEY)

def list_folders(bucket_name):
    try:
        folders = []
        result = s3.list_objects_v2(Bucket=bucket_name, Delimiter='/')
        for o in result.get('CommonPrefixes', []):
            folders.append(o.get('Prefix'))
        return folders
    except NoCredentialsError:
        st.error("Credentials not available")
        return []
    except ClientError as e:
        st.error(f"Client error: {e}")
        return []
    except EndpointConnectionError as e:
        st.error(f"Endpoint connection error: {e}")
        return []

def count_files_in_folder(bucket_name, prefix):
    try:
        folder_counts = defaultdict(lambda: defaultdict(int))
        paginator = s3.get_paginator('list_objects_v2')
        pages = paginator.paginate(Bucket=bucket_name, Prefix=prefix)
        for page in pages:
            for obj in page.get('Contents', []):
                folder = '/'.join(obj['Key'].split('/')[:-1])
                last_modified = obj['LastModified'].strftime('%Y-%m-%d')
                folder_counts[folder][last_modified] += 1
        return folder_counts
    except NoCredentialsError:
        st.error("Credentials not available")
        return {}
    except ClientError as e:
        st.error(f"Client error: {e}")
        return {}
    except EndpointConnectionError as e:
        st.error(f"Endpoint connection error: {e}")
        return {}

# Streamlit UI
st.title('NeuroSinQ IAT Completions Dashboard')

# Dropdown to select folder
folders = list_folders(BUCKET_NAME)
#selected_folder = st.selectbox('Select a folder', folders)
selected_folder = st.selectbox('Select a folder', folders, disabled=True, index=folders.index('JUNIPER/') if 'JUNIPER/' in folders else 0)
# Display file counts
if selected_folder:
    folder_counts = count_files_in_folder(BUCKET_NAME, selected_folder)
    
    if folder_counts:
        # Collect all dates
        all_dates = sorted(set(date for counts in folder_counts.values() for date in counts))
        
        # Prepare data for display
        data = []
        for folder, counts in folder_counts.items():
            row = {'Folder': folder}
            for date in all_dates:
                row[date] = counts.get(date, 0)
            data.append(row)
        
        st.write("File counts by folder and last modified date:")
        st.table(data)