import streamlit as st
from unsloth import FastVisionModel
from transformers import AutoModel
from PIL import Image
import base64
import os
import tempfile

os.environ["UNSLOTH_WARN_UNINITIALIZED"] = '0'

# Page configuration
st.set_page_config(
    page_title="DeepSeek-OCR",
    page_icon="🔎",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Load model with caching (only loads once)
@st.cache_resource
def load_model():
    """Load the DeepSeek OCR model once and cache it."""
    model, tokenizer = FastVisionModel.from_pretrained(
    model_name="Baldezo313/deepseek_ocr_merged",
    load_in_4bit=True,  # recommandé sur HF Spaces
    auto_model=AutoModel,
    trust_remote_code=True,
    )
    FastVisionModel.for_inference(model)  # Enable for inference!
    
    return model, tokenizer

# Title and description in main area
logo_path = "./assets/deepseek.png"

if os.path.exists(logo_path):
    st.markdown(
        """
        # <img src="data:image/png;base64,{}" width="50" style="vertical-align: -12px;"> DeepSeek-OCR
        """.format(
            base64.b64encode(open(logo_path, "rb").read()).decode()
        ),
        unsafe_allow_html=True,
    )
else:
    st.title("DeepSeek-OCR")

# Add clear button to top right
col1, col2 = st.columns([6,1])
with col2:
    if st.button("Clear 🗑️"):
        if 'ocr_result' in st.session_state:
            del st.session_state['ocr_result']
        st.rerun()

st.markdown('<p style="margin-top: -20px;">Extract structured text from images using DeepSeek-OCR!</p>', unsafe_allow_html=True)
st.markdown("---")

# Move upload controls to sidebar
with st.sidebar:
    st.header("Upload Image")
    uploaded_file = st.file_uploader("Choose an image...", type=['png', 'jpg', 'jpeg'])
    
    if uploaded_file is not None:
        # Display the uploaded image
        image = Image.open(uploaded_file)
        st.image(image, caption="Uploaded Image")
        
        if st.button("Extract Text 🔍", type="primary"):
            with st.spinner("Processing image..."):
                try:
                    # Load model (cached, so only loads once)
                    # Show loading message if model hasn't been loaded yet
                    if 'model_loaded' not in st.session_state:
                        st.info("Loading DeepSeek OCR model (this only happens once)...")
                        st.session_state['model_loaded'] = True
                    model, tokenizer = load_model()
                    
                    # Save uploaded file temporarily (preserve original extension)
                    file_extension = os.path.splitext(uploaded_file.name)[1] or '.jpg'
                    with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as tmp_file:
                        tmp_file.write(uploaded_file.getvalue())
                        tmp_image_path = tmp_file.name
                    
                    # Create temporary output directory
                    with tempfile.TemporaryDirectory() as tmp_output_dir:
                        # Run inference
                        prompt = "<image>\nFree OCR. "
                        res = model.infer(
                            tokenizer,
                            prompt=prompt,
                            image_file=tmp_image_path,
                            output_path=tmp_output_dir,
                            image_size=640,
                            base_size=1024,
                            crop_mode=True,
                            save_results=True,
                            test_compress=False
                        )
                        
                        # Try to read from output file if result is saved there
                        result_file = os.path.join(tmp_output_dir, 'result.mmd')
                        if os.path.exists(result_file):
                            with open(result_file, 'r', encoding='utf-8') as f:
                                ocr_result = f.read()
                        else:
                            ocr_result = "No text found in the image."
                        
                        st.session_state['ocr_result'] = ocr_result
                    
                    # Clean up temporary image file
                    os.unlink(tmp_image_path)
                    
                except Exception as e:
                    st.error(f"Error processing image: {str(e)}")
                    import traceback
                    st.error(f"Details: {traceback.format_exc()}")

# Main content area for results
if 'ocr_result' in st.session_state:
    st.text_area(
    "OCR Result",
    st.session_state['ocr_result'],
    height=500
    )
else:
    st.info("Upload an image and click 'Extract Text' to see the results here.")

# Footer
st.markdown("---")
st.markdown("Made with ❤️ using [DeepSeek-OCR](https://huggingface.co/deepseek-ai/DeepSeek-OCR) Model | [Report an Issue](https://github.com/patchy631/ai-engineering-hub/issues)")