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
from models.text_processor import process_menu_text

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

# Sidebar for model settings
st.sidebar.header("Settings")
use_llm = st.sidebar.checkbox("Use LLM for text processing", value=True)

# Add information about the application
st.sidebar.markdown("---")
st.sidebar.subheader("About")
st.sidebar.info(
    "This application converts menu images to Braille text using AI. "
    "It extracts text from images using document AI, processes the text with LLMs, "
    "and will convert to Braille in future versions."
)

# 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("Process Menu"):
        with st.spinner("Processing image..."):
            # Preprocess the image
            preprocessed_img = preprocess_image(image)
            
            # Extract text using LayoutLMv2
            try:
                result = extract_text_and_layout(preprocessed_img)
                
                # Display extracted words
                if result['words']:
                    raw_text = ' '.join(result['words'])
                    
                    # Show raw text in an expandable section
                    with st.expander("Raw Extracted Text"):
                        st.text_area("Raw OCR Output", raw_text, height=150)
                    
                    # Process text with LLM if enabled
                    if use_llm:
                        st.subheader("Processed Menu Text")
                        with st.spinner("Enhancing text with AI..."):
                            processed_result = process_menu_text(raw_text)
                            
                            if processed_result['success']:
                                st.text_area("Structured Menu Text", 
                                           processed_result['structured_text'], 
                                           height=300)
                                
                                # Store the processed result for later use
                                st.session_state.processed_text = processed_result['structured_text']
                                st.session_state.menu_data = processed_result.get('menu_data', {})
                            else:
                                st.warning(f"AI processing failed: {processed_result.get('error', 'Unknown error')}")
                                st.text_area("Text Output", raw_text, height=300)
                                st.session_state.processed_text = raw_text
                    else:
                        # Just use the raw text
                        st.subheader("Extracted Text")
                        st.text_area("Text Output", raw_text, height=300)
                        st.session_state.processed_text = raw_text
                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")