# app.py (개별 네비게이션 및 선택 드롭다운 동시 지원)
import streamlit as st
import json
from pathlib import Path
import base64

# --- 헬퍼 함수 ---

def match_uploaded_files(doc_files, json_files):
    """
    업로드된 두 파일 목록을 받아, 이름(확장자 제외)을 기준으로 매칭하고
    결과를 딕셔너리로 반환합니다.
    """
    matched_pairs = {}
    docs_map = {Path(f.name).stem: f for f in doc_files}
    jsons_map = {Path(f.name).stem: f for f in json_files}

    for stem, doc_file in docs_map.items():
        if stem in jsons_map:
            matched_pairs[stem] = {
                "doc_file": doc_file,
                "json_file": jsons_map[stem]
            }
    return matched_pairs

def display_pdf(file_object):
    """
    업로드된 파일 객체(UploadedFile)를 읽어 PDF를 표시합니다.
    """
    try:
        file_object.seek(0)
        base64_pdf = base64.b64encode(file_object.read()).decode('utf-8')
        pdf_display = f'<iframe src="data:application/pdf;base64,{base64_pdf}" width="100%" height="800" type="application/pdf"></iframe>'
        st.markdown(pdf_display, unsafe_allow_html=True)
    except Exception as e:
        st.error(f"PDF 파일을 표시하는 중 오류가 발생했습니다: {e}")

# --- 콜백 함수 ---

def handle_nav_button(direction, total_files):
    """이전/다음 버튼 클릭을 처리하는 콜백"""
    if direction == "prev" and st.session_state.current_index > 0:
        st.session_state.current_index -= 1
    elif direction == "next" and st.session_state.current_index < total_files - 1:
        st.session_state.current_index += 1

def handle_selectbox_change():
    """selectbox 변경을 처리하는 콜백"""
    selected_basename_with_index = st.session_state.selectbox_key
    # "1. filename" 형식에서 인덱스만 추출
    new_index = int(selected_basename_with_index.split('. ', 1)[0]) - 1
    st.session_state.current_index = new_index

# --- 메인 UI 로직 ---

def main():
    st.set_page_config(layout="wide", page_title="결과 비교 도구")
    st.title("📑 파일 업로드 기반 결과 비교 도구")
    st.markdown("---")

    # --- 1. 파일 업로드 ---
    st.sidebar.header("파일 업로드")
    
    uploaded_docs = st.sidebar.file_uploader(
        "1. 원본 문서 파일(들)을 업로드하세요.",
        accept_multiple_files=True,
        type=['png', 'jpg', 'jpeg', 'pdf']
    )
    
    uploaded_jsons = st.sidebar.file_uploader(
        "2. 결과 JSON 파일(들)을 업로드하세요.",
        accept_multiple_files=True,
        type=['json']
    )

    if not uploaded_docs or not uploaded_jsons:
        st.info("사이드바에서 원본 문서와 결과 JSON 파일을 모두 업로드해주세요.")
        return

    try:
        matched_files = match_uploaded_files(uploaded_docs, uploaded_jsons)
    except Exception as e:
        st.error(f"업로드된 파일을 매칭하는 중 오류가 발생했습니다: {e}")
        return

    if not matched_files:
        st.warning("업로드된 파일 중 일치하는 문서-JSON 쌍을 찾을 수 없습니다. 파일 이름(확장자 제외)이 동일한지 확인하세요.")
        return

    # --- 2. 파일 네비게이션 및 선택 ---
    st.sidebar.header("파일 탐색")
    
    sorted_basenames = sorted(list(matched_files.keys()))
    total_files = len(sorted_basenames)

    # 세션 상태 초기화
    if 'current_index' not in st.session_state:
        st.session_state.current_index = 0
    
    # 인덱스가 유효한 범위를 벗어나지 않도록 조정
    st.session_state.current_index = max(0, min(st.session_state.current_index, total_files - 1))

    # Selectbox
    display_options = [f"{i + 1}. {name}" for i, name in enumerate(sorted_basenames)]
    st.selectbox(
        "파일을 직접 선택하세요:",
        options=display_options,
        index=st.session_state.current_index,
        key='selectbox_key', # 콜백에서 값을 참조하기 위한 키
        on_change=handle_selectbox_change
    )

    # 네비게이션 컨트롤
    col1, col2, col3 = st.sidebar.columns([1, 2, 1])
    
    col1.button(
        "◀ 이전", 
        on_click=handle_nav_button, 
        args=("prev", total_files), 
        use_container_width=True
    )
    col2.markdown(f"<p style='text-align: center;'>{st.session_state.current_index + 1} / {total_files}</p>", unsafe_allow_html=True)
    col3.button(
        "다음 ▶", 
        on_click=handle_nav_button, 
        args=("next", total_files), 
        use_container_width=True
    )

    # --- 3. 결과 표시 ---
    current_basename = sorted_basenames[st.session_state.current_index]
    st.header(f"🔎 비교 결과: `{current_basename}`")
    
    selected_pair = matched_files[current_basename]
    doc_file = selected_pair["doc_file"]
    json_file = selected_pair["json_file"]

    res_col1, res_col2 = st.columns(2)
    with res_col1:
        st.subheader(f"원본 문서: `{doc_file.name}`")
        doc_suffix = Path(doc_file.name).suffix.lower()
        
        if doc_suffix == ".pdf":
            display_pdf(doc_file)
        else:
            st.image(doc_file, caption=f"원본 이미지: {doc_file.name}", use_container_width=True)

    with res_col2:
        st.subheader(f"추출된 데이터: `{json_file.name}`")
        try:
            json_file.seek(0)
            data = json.load(json_file)
            
            result_to_display = data[0] if isinstance(data, list) and data else data
            
            if isinstance(result_to_display, dict) and 'fields' in result_to_display:
                del result_to_display['fields']
                
            st.json(result_to_display)
        except Exception as e:
            st.error(f"JSON 파일을 읽거나 처리하는 중 오류가 발생했습니다: {e}")

if __name__ == "__main__":
    main()