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

# AWS S3 setup
aws_access_key=os.environ['AWS_ACCESS_KEY']
aws_secret_key=os.environ['AWS_SECRET_KEY']
aws_bucket=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(aws_bucket)
selected_folder = st.selectbox('Select a folder', folders, index=folders.index('/') if '/' in folders else 0)

# Display file counts
if selected_folder:
    folder_counts = count_files_in_folder(aws_bucket, 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 = []
        totals = defaultdict(int)
        grand_total = 0
        for folder, counts in folder_counts.items():
            row = {'Folder': folder}
            row_total = 0
            for date in all_dates:
                row[date] = counts.get(date, 0)
                totals[date] += counts.get(date, 0)
                grand_total += counts.get(date, 0)
                row_total += counts.get(date, 0)
            row['Total'] = row_total
            data.append(row)
        
        # Add totals row
        total_row = {'Folder': 'TOTAL'}
        for date in all_dates:
            total_row[date] = totals[date]
        total_row['Total'] = grand_total
        data.append(total_row)
        
        st.write("File counts by folder and last modified date:")
        st.table(data)