import streamlit as st
import pandas as pd
from io import BytesIO
from toshiba import extract_toshiba_table
from bhel import extract_bhel_table
from federal_electric import extract_federal_electric_table
from al_nisf import extract_al_nisf_table
from others import extract_others_table

# Streamlit app
def main():
    st.title("PO Data Processor")
    st.markdown("### Extract and Format PO Data")

    # Dropdown for PO type
    po_options = ["Toshiba", "BHEL", "Federal Electric", "AL NISF", "Others"]
    selected_po_type = st.selectbox("Select a PO type:", po_options)

    # File upload
    uploaded_file = st.file_uploader("Upload your PO PDF file:", type=["pdf"])
    if not uploaded_file:
        st.info("Please upload a PDF file to proceed.")
        return

    # Process based on selected PO type
    st.write(f"Processing PO type: {selected_po_type}")
    try:
   iif selected_po_type == "Toshiba":
    df = extract_toshiba_table(uploaded_file)  # Indented 4 spaces
elif selected_po_type == "BHEL":
    df = extract_bhel_table(uploaded_file)  # Indented 4 spaces
elif selected_po_type == "Federal Electric":
    df = extract_federal_electric_table(uploaded_file)  # Indented 4 spaces
elif selected_po_type == "AL NISF":
    df = extract_al_nisf_table(uploaded_file)  # Indented 4 spaces
else:
    st.warning("Other PO types are not yet implemented. Please use a supported format.")
    return  # Indented 4 spaces


        if df.empty:
            st.warning("No tabular data found in the uploaded PDF.")
            return
    except Exception as e:
        st.error(f"Error processing the PDF: {e}")
        return

    # Display the extracted data
    st.write("### Extracted Data:")
    st.dataframe(df)

    # Export options
    # CSV Export
    csv_data = df.to_csv(index=False).encode("utf-8")
    st.download_button(
        label="Download as CSV",
        data=csv_data,
        file_name=f"{selected_po_type.lower()}_po_data.csv",
        mime="text/csv"
    )

    # Excel Export
    output = BytesIO()
    with pd.ExcelWriter(output, engine='openpyxl') as writer:
        df.to_excel(writer, index=False, sheet_name="Sheet1")
    output.seek(0)  # Reset buffer pointer
    excel_data = output.getvalue()

    st.download_button(
        label="Download as Excel",
        data=excel_data,
        file_name=f"{selected_po_type.lower()}_po_data.xlsx",
        mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
    )

if __name__ == "__main__":
    main()