import streamlit as st
import openai
import pandas as pd
import datetime

# Initialize Streamlit app
st.set_page_config(page_title="OpenAI Usage Dashboard", layout="centered")

# Set up OpenAI API key
st.sidebar.title("OpenAI API Usage Dashboard")
api_key = st.sidebar.text_input("Enter your OpenAI API Key:", type="password")
openai.api_key = api_key

# Title and description
st.title("OpenAI Usage Tracker")
st.write("This app retrieves and displays your OpenAI API usage data for a specified date range.")

# Date selection
start_date = st.sidebar.date_input("Select Start Date:", datetime.date.today() - datetime.timedelta(days=30))
end_date = st.sidebar.date_input("Select End Date:", datetime.date.today())

# Ensure end date is not before start date
if start_date > end_date:
    st.error("End Date must be after Start Date")

# Function to retrieve usage data
def get_usage_data(api_key, start_date, end_date):
    try:
        openai.api_key = api_key
        usage = openai.api_usage.get(start_date=start_date, end_date=end_date)
        usage_data = usage['data']
        # Create a DataFrame from the usage data
        data = pd.DataFrame(usage_data)
        data['date'] = pd.to_datetime(data['date'])
        data.set_index('date', inplace=True)
        return data[['n_tokens_total']]
    except Exception as e:
        st.error(f"An error occurred: {e}")
        return None

# Retrieve and display usage data if API key is provided
if api_key:
    usage_data = get_usage_data(api_key, start_date.strftime("%Y-%m-%d"), end_date.strftime("%Y-%m-%d"))

    # Display usage data as a table and line chart
    if usage_data is not None:
        st.subheader("Usage Data")
        st.write("Displaying usage from", start_date, "to", end_date)
        st.dataframe(usage_data)

        # Display line chart for usage over time
        st.subheader("Token Usage Over Time")
        st.line_chart(usage_data['n_tokens_total'])

        # Display statistics
        total_usage = usage_data['n_tokens_total'].sum()
        avg_usage = usage_data['n_tokens_total'].mean()
        st.subheader("Statistics")
        st.write(f"Total Token Usage: {total_usage}")
        st.write(f"Average Daily Token Usage: {avg_usage:.2f}")
    else:
        st.info("Enter your API key to view usage data.")
else:
    st.warning("Please enter your OpenAI API key to proceed.")