import logging 
import streamlit as st
import pandas as pd
import io
from PIL import Image, ImageDraw, ImageFont
from myocr.pipelines import CommonOCRPipeline

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
)
logger = logging.getLogger(__name__)

st.set_page_config(layout="wide")
st.title("MyOCR Demo")

# init pipeline
@st.cache_resource
def load_pipeline():
    return CommonOCRPipeline("cpu")

pipeline = load_pipeline()

font = ImageFont.truetype("src/NotoSans.ttf", 12)

def process_image(image: Image.Image, format):
    buffer = io.BytesIO()
    image.save(buffer, format=format)
    results = pipeline(buffer.getvalue())
    logger.info(f"ocr results: {results}")
    if not results or not results.regions:
        return None, []

    image_with_boxes = image.copy()
    draw = ImageDraw.Draw(image_with_boxes)
    table_data = []

    for item in results.regions:
        if item.confidence <= 0.5:
            continue
        shape = item.bounding_shape
        points = shape.points
        if shape.type == "rectangle":
            draw.rectangle([(points[0].x, points[0].y), (points[2].x, points[2].y)], outline="red", width=2)
            text_pos_y = max(points[0].y - 18, 0)
            draw.text((points[0].x, text_pos_y), item.text, font=font, fill="green")
        table_data.append((item.text, item.confidence))

    return image_with_boxes, table_data 

def main():
    left_col, right_col = st.columns([2, 1.5])
    with left_col:
        uploaded_file = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg"])
        if not uploaded_file:
            return
        try:
            image = Image.open(uploaded_file).convert("RGB")
        except Exception:
            st.error("⚠️ Invalid image file. Please upload a valid PNG or JPG.")
            logger.warning("Invalid file uploaded.")
            return
        
        spinner_container = st.empty()
        image_slot = st.empty()
        image_slot.image(image, use_container_width=True)

        with spinner_container:
            with st.spinner("Recognizing text..."):
                mime_type = uploaded_file.type
                processed_image, table_data = process_image(image, mime_type.split("/")[-1])

        if not table_data:
            right_col.warning("No text detected.")
            return
        image_slot.image(processed_image, use_container_width=True)

    # --- Show results ---
    with right_col:
        st.subheader("Recognized Text")
        df = pd.DataFrame(table_data, columns=["Text", "Confidence"])
        st.table(df)

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        logger.exception("Unhandled exception occurred.")
        st.error(f"Internal Error!")