import streamlit as st
from PIL import Image
import pandas as pd
import numpy as np
from invoice import extract_data, extract_tables, INVOICE




def process_image(lang, to_be_extracted, input_image):
    data = extract_data(lang, to_be_extracted, input_image)
    return data

def main():
    st.title("Image Processing App")
    st.write("Upload an image and click the 'Process Image' button to process it.")
    uploaded_image = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png","webp"])
    
    if uploaded_image is not None:
        # Display the uploaded image
        st.image(uploaded_image, caption="Uploaded Image", use_column_width=True)

        lang = st.selectbox("Select Language", ["french", "english", "arabic"])
        pil_image = Image.open(uploaded_image).convert('RGB')
        numpy_image = np.array(pil_image)

        st.write("Add elements to extract:")
        extract_input = st.text_input("Add elements")
        extract_list = st.session_state.get("extract_list", INVOICE)
        if extract_input:
            extract_list.append(extract_input.strip())
            st.session_state["extract_list"] = extract_list

        # Display the extract list as chips
        st.write("Elements to extract:")
        for item in extract_list:
            st.write(f"`{item}`", unsafe_allow_html=True)

        if st.button("Extract information"):
            
            image_info = process_image(lang, extract_list, numpy_image)

            df = pd.DataFrame(list(image_info.items()), columns=["Field", "Value"])
            st.write("Extracted information:")
            st.dataframe(df)
    
        if st.button("Extract Tables"):
            csv = extract_tables(lang, numpy_image)
            st.download_button(label="Download CSV", data=csv, file_name='data.csv', mime='text/csv')
        
            
    else:
        st.session_state['extract_list'] = INVOICE

if __name__ == "__main__":
    main()