import streamlit as st
from PIL import Image
import io
import numpy as np

# Import our custom modules
from utils.image_preprocessing import preprocess_image
from models.document_ai import extract_text_and_layout

# App title and description
st.title("Menu to Braille Converter")
st.write("Upload a menu image to convert it to Braille text")

# File uploader
uploaded_file = st.file_uploader("Choose a menu image...", type=["jpg", "jpeg", "png"])

# Display uploaded image and process it
if uploaded_file is not None:
    # Load and display image
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Menu", use_column_width=True)
    
    # Add a button to process the image
    if st.button("Extract Text"):
        with st.spinner("Processing image..."):
            # Preprocess the image
            st.subheader("Preprocessed Image")
            preprocessed_img = preprocess_image(image)
            st.image(preprocessed_img, caption="Preprocessed Image", use_column_width=True)
            
            # Extract text using LayoutLMv2
            st.subheader("Extracted Text")
            try:
                result = extract_text_and_layout(preprocessed_img)
                
                # Display extracted words
                if result['words']:
                    text = ' '.join(result['words'])
                    st.text_area("Extracted Text", text, height=200)
                else:
                    st.warning("No text was extracted from the image.")
                    
            except Exception as e:
                st.error(f"Error processing image: {str(e)}")
    
    # Placeholders for future functionality
    st.subheader("Braille Translation")
    st.info("Braille translation will be implemented in Phase 4")
    
    st.subheader("Download Options")
    st.info("PDF download will be implemented in Phase 5")