import streamlit as st
import wandb
import pandas as pd
import os
import matplotlib.pyplot as plt

from utils import fetch_runs_to_df

# Access the API key from the environment variable
wandb_api_key = os.getenv('WANDB_API_KEY')

# Log in to wandb using the API key
if wandb_api_key:
    wandb.login(key=wandb_api_key)
else:
    st.error("WANDB_API_KEY not found in environment variables.")

# Initialize W&B API
api = wandb.Api()


# Define available projects (bookmarks)
projects = {
    "Competition 1": {"entity": "urbaniak-bruno-safescanai", "project": "pytorch-intro"},
    "Competition 2": {"entity": "urbaniak-bruno-safescanai", "project": "basic-intro"},
    # Add more projects as needed
}
bookmarks = ["All", "Competition 1", "Competition 2"]

# Sidebar for project selection
st.sidebar.title("Bookmarks")
selected_project = st.sidebar.selectbox("Select a competition:", bookmarks)


# Sidebar with buttons
st.sidebar.title("Iluzja wyboru")
option = st.sidebar.radio(
    "Select an option:",
    ["General", "Researchers", "Validators", "Models"]
)

df = fetch_runs_to_df(api, projects, selected_project)

# Streamlit UI
st.title("W&B Data in Streamlit")

# Sidebar filter options
st.sidebar.header("Filter Options")
run_name_filter = st.sidebar.text_input("Filter by Run Name", "")
state_filter = st.sidebar.selectbox("Filter by State", ["All"] + df["State"].unique().tolist())
tag_filter = st.sidebar.text_input("Filter by Tags (comma separated)", "")

# Apply filters
if run_name_filter:
    df = df[df["Run Name"].str.contains(run_name_filter, case=False, na=False)]
if state_filter != "All":
    df = df[df["State"] == state_filter]
if tag_filter:
    tag_list = [tag.strip() for tag in tag_filter.split(",")]
    df = df[df["Tags"].apply(lambda tags: any(tag in tags for tag in tag_list))]

# Display the filtered DataFrame
st.dataframe(df)

# Display details of selected run
selected_run_id = st.selectbox("Select a Run ID to see details", df["ID"].tolist() if not df.empty else [])
if selected_run_id:
    selected_run = api.run(f"{projects[selected_project]['entity']}/{projects[selected_project]['project']}/{selected_run_id}")
    run_df = selected_run.history()
    st.write(f"Details for run: {selected_run.name}")
    st.dataframe(run_df)