import streamlit as st
import pandas as pd
import io
from utils import get_db_connection, fetch_widgets, get_s3_client, S3_BUCKET

st.header("Add New Widget & View Table")

# Load widgets data (cached in session state so we query only once per page load)
if "widgets_df" not in st.session_state:
    try:
        st.session_state.widgets_df = fetch_widgets()
    except Exception as e:
        st.error(f"Error fetching data: {e}")

st.subheader("Widgets Table")
st.dataframe(st.session_state.widgets_df)

st.subheader("Add New Widget")

with st.form("widget_form", clear_on_submit=True):
    col1, col2, col3 = st.columns(3)
    with col1:
        height = st.number_input("Height", min_value=0.0, format="%.2f", step=0.1)
        width = st.number_input("Width", min_value=0.0, format="%.2f", step=0.1)
    with col2:
        length = st.number_input("Length", min_value=0.0, format="%.2f", step=0.1)
        quantity = st.number_input("Quantity", min_value=1, step=1)
    with col3:
        color = st.selectbox("Color", options=["red", "blue", "green"])
        notes = st.text_area("Notes")
    # Marking the upload as optional in the label
    file = st.file_uploader("Upload File (Optional)", type=["txt", "csv", "png", "jpg", "jpeg", "pdf"])
    submitted = st.form_submit_button("Submit")

if submitted:
    try:
        # Default S3 key to None if no file is provided.
        s3_key = None
        if file is not None:
            # Upload file to S3 using the original filename as the object key
            s3 = get_s3_client()
            s3.upload_fileobj(file, Bucket=S3_BUCKET, Key=file.name)
            s3_key = file.name

        # Insert the new widget into PostgreSQL
        conn = get_db_connection()
        cur = conn.cursor()
        insert_query = """
            INSERT INTO widgets (height, width, length, color, quantity, notes, s3_key)
            VALUES (%s, %s, %s, %s, %s, %s, %s)
        """
        cur.execute(insert_query, (height, width, length, color, quantity, notes, s3_key))
        conn.commit()
        cur.close()
        conn.close()

        st.success("Widget added successfully!")

        # Refresh the widgets table
        st.session_state.widgets_df = fetch_widgets()
        st.dataframe(st.session_state.widgets_df)
    except Exception as e:
        st.error(f"An error occurred: {e}")