import streamlit as st
import cv2
import pytesseract
import numpy as np
import pandas as pd
from extractor import extract_multi_records

st.set_page_config(page_title="Smart Multi-Document Data Extractor", layout="wide")

st.title("📄 Smart Multi-Document Data Extractor")

# OCR Function (Improved)
def extract_text(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gray = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY)[1]
    return pytesseract.image_to_string(gray)

# File uploader (multiple files supported)
uploaded_files = st.file_uploader(
    "Upload Images",
    type=["jpg", "png", "jpeg"],
    accept_multiple_files=True
)

if uploaded_files:

    all_data = []

    for uploaded_file in uploaded_files:
        try:
            file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
            img = cv2.imdecode(file_bytes, 1)

            # Check image valid
            if img is None:
                st.error(f"❌ Could not read image: {uploaded_file.name}")
                continue

            # Show image
            st.image(img, caption=uploaded_file.name, use_container_width=True)

            # Extract text
            text = extract_text(img)

            if not text.strip():
                st.warning(f"⚠️ No text found in {uploaded_file.name}")
                continue

            # Extract structured data
            records = extract_multi_records(text)

            if not records:
                st.warning(f"⚠️ No structured data found in {uploaded_file.name}")
                continue

            all_data.extend(records)

        except Exception as e:
            st.error(f"❌ Error processing {uploaded_file.name}: {str(e)}")

    # Show final data
    if all_data:
        df = pd.DataFrame(all_data)

        st.success("✅ Extraction Complete!")
        st.dataframe(df)

        # Download Excel
        df.to_excel("output.xlsx", index=False)

        with open("output.xlsx", "rb") as f:
            st.download_button(
                "📥 Download Excel",
                f,
                file_name="output.xlsx"
            )
    else:
        st.warning("⚠️ No data extracted from uploaded files.")