import streamlit as st
from pathlib import Path
from paddleocr import PPStructureV3
from PIL import Image
import json
import os
import uuid
import numpy as np

# 페이지 설정
st.set_page_config(page_title="PaddleOCR 이미지 분석", layout="wide")

st.title("📄 PaddleOCR을 이용한 이미지 분석")
st.write("이미지 파일을 업로드하면 문서 방향 분류, 왜곡 보정, 표/수식/차트 인식 등을 수행하고 결과를 보여줍니다.")

# 모델 초기화 (캐싱 사용)
@st.cache_resource
def load_model():
    return PPStructureV3(
        lang="korean",
        use_doc_orientation_classify=True, 
        use_doc_unwarping=True,
        use_seal_recognition=False,
        use_table_recognition=True,
        use_formula_recognition=True,
        use_chart_recognition=True,
        use_region_detection=True,
    )

with st.spinner("모델을 불러오는 중입니다..."):
    structure = load_model()

# 파일 업로더
uploaded_file = st.file_uploader("분석할 이미지 파일을 선택하세요.", type=["jpg", "jpeg", "png", "bmp", "tif"])

if uploaded_file is not None:
    # 임시 디렉터리 설정
    output_dir = Path("output_results")
    output_dir.mkdir(exist_ok=True)
    
    # 고유한 파일 이름 생성
    unique_id = uuid.uuid4().hex
    
    try:
        # 원본 이미지 표시
        image = Image.open(uploaded_file)
        
        # 모델이 처리하기 쉽도록 이미지를 RGB 형식으로 변환
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        col1, col2 = st.columns(2)
        with col1:
            st.subheader("🖼️ 원본 이미지")
            st.image(image, caption="업로드된 이미지", width='stretch')

        # 분석 시작 버튼
        if st.button("분석 시작하기", use_container_width=True):
            with st.spinner("이미지를 분석하고 있습니다..."):
                # PIL 이미지를 numpy 배열로 변환하여 예측
                img_array = np.array(image)
                output = structure.predict(input=img_array)

                if not output:
                    st.warning("이미지에서 구조를 감지하지 못했습니다.")
                else:
                    # 결과 저장 경로
                    saved_res_path_base = output_dir / f"result_{unique_id}"
                    saved_res_path_base.mkdir(exist_ok=True)
                    
                    json_paths = []
                    # 1. 모든 결과 파일(이미지, JSON)을 먼저 저장
                    for i, res in enumerate(output):
                        res.save_to_img(save_path=str(saved_res_path_base))
                        
                        json_path = saved_res_path_base / f"{i}.json"
                        res.save_to_json(save_path=str(json_path))
                        json_paths.append(json_path)

                    # 2. 저장된 모든 결과 이미지를 설명과 함께 col2에 표시
                    with col2:
                        st.subheader("✨ 분석 결과 이미지")
                        image_files = sorted([f for f in saved_res_path_base.glob('*') if f.suffix.lower() in ('.jpg', '.jpeg', '.png', '.bmp', '.tif')])
                        
                        if not image_files:
                            st.error(f"결과 이미지 파일을 찾을 수 없습니다. (검색 경로: {saved_res_path_base})")
                        else:
                            for img_path in image_files:
                                result_image = Image.open(img_path)
                                
                                title = ""
                                description = ""

                                if "_layout_order_res" in img_path.name:
                                    title = "레이아웃 순서 분석 (Reading Order)"
                                    description = "각 텍스트 영역을 사람이 문서를 읽는 논리적인 순서(예: 위에서 아래로, 왼쪽에서 오른쪽으로)를 파악하여 시각화한 결과입니다. 복잡한 문서에서 텍스트를 올바른 순서로 추출하는 데 중요한 역할을 합니다."
                                elif "layout_det_res" in img_path.name:
                                    title = "레이아웃 감지 (Layout Detection)"
                                    description = "모델이 문서에서 텍스트, 표, 이미지 등의 영역을 최초로 감지한 결과입니다. 이 결과를 바탕으로 각 영역의 종류를 더 상세하게 분석합니다."
                                elif "_region_det_res" in img_path.name:
                                    title = "영역 감지 (Region Detection)"
                                    description = "레이아웃 감지 후, 각 영역의 종류('제목', '표', '텍스트' 등)를 구체적으로 식별해낸 결과입니다."
                                elif "_overall_ocr_res" in img_path.name:
                                    title = "종합 OCR 결과 (Overall OCR Result)"
                                    description = "레이아웃 분석으로 찾아낸 각 텍스트 영역에 대해 광학 문자 인식(OCR)을 수행한 최종 결과를 원본 이미지 위에 표시한 것입니다."
                                elif "_preprocessed_img" in img_path.name:
                                    title = "전처리된 이미지 (Preprocessed Image)"
                                    description = "OCR의 정확도를 높이기 위해 입력 이미지의 기울기를 보정하거나, 잡음을 제거하는 등의 전처리 과정을 거친 후의 이미지를 보여줍니다."
                                elif "_table_cell_img" in img_path.name:
                                    title = "표 셀 인식 (Table Cell Recognition)"
                                    description = "'표(Table)' 영역이 감지되면, 표의 구조를 분석하여 각 셀의 경계를 찾고 그 결과를 시각화하여 보여줍니다. 이 결과를 바탕으로 구조화된 데이터를 생성할 수 있습니다."
                                else:
                                    title = "기타 결과 이미지"
                                    description = "분석 과정에서 생성된 기타 결과 이미지입니다."

                                st.markdown(f"##### {title}")
                                st.image(result_image, width='stretch')
                                st.info(description)
                                st.markdown("---")

                    # 3. 저장된 모든 JSON 결과를 아래에 순서대로 표시
                    st.subheader("📄 분석 내용")
                    for json_path in sorted(json_paths):
                        with open(json_path, 'r', encoding='utf-8') as f:
                            json_data = json.load(f)
                        
                        st.markdown(f"---")
                        st.markdown(f"#### 결과 파일: `{json_path.name}`")

                        # 테이블 HTML을 먼저 추출하여 저장
                        table_htmls = [table.get('pred_html', '') for table in json_data.get('table_res_list', [])]
                        table_idx = 0

                        # 각 블록을 순회하며 의미에 맞게 표시
                        for block in json_data.get('parsing_res_list', []):
                            label = block.get('block_label', 'unknown')
                            content = block.get('block_content', '').strip()

                            if label == 'table':
                                st.markdown(f"##### 📊 표 (Table)")
                                if table_idx < len(table_htmls):
                                    st.markdown(table_htmls[table_idx], unsafe_allow_html=True)
                                    table_idx += 1
                                else:
                                    st.warning("테이블 내용은 찾을 수 없습니다.")
                            
                            elif label in ['header', 'doc_title']:
                                st.markdown(f"##### 📑 제목 / 헤더 ({label})")
                                st.markdown(f"**{content}**")
                            
                            elif label in ['image', 'seal']:
                                st.markdown(f"##### 🖼️ 이미지 / 직인 ({label})")
                                st.info(f"'{label}' 영역이 감지되었습니다.")
                            
                            elif content:
                                st.markdown(f"##### 📝 텍스트 ({label})")
                                st.write(content)

                        with st.expander(f"`{json_path.name}` 전체 내용 보기"):
                            st.json(json_data)

                # 임시 파일 삭제 (이제 필요 없음)
                # if temp_image_path.exists():
                #     os.remove(temp_image_path)

    except Exception as e:
        st.error(f"이미지 처리 중 예상치 못한 오류가 발생했습니다: {e}")
        st.error(f"파일: {uploaded_file.name}")
        st.info("이미지 파일이 손상되었거나, 모델이 지원하지 않는 형식일 수 있습니다. 다른 이미지로 시도해 보세요.")

else:
    st.info("분석할 이미지를 업로드해주세요.")