import streamlit as st
from app_logic import load_data, entities2Selection
from app_ui import (
    setup_report_selection,
    display_report_content,
    display_entities,
    display_relationship_graph,
    handle_review_submission,
    setup_input_selection
)

def initialize_session_state():
    """Initialize session state variables"""
    if 'reports_json' not in st.session_state:
        
        file_path = 'mockedReports.json'

        st.session_state.reports_json = load_data(file_path)
    if 'selected_entity' not in st.session_state:
        st.session_state.selected_entity = None

def main():
    """Main application"""
    # 设置页面为 wide mode
    st.set_page_config(
        page_title="Medical Report Review System",
        layout="wide",
        initial_sidebar_state="expanded"
    )

    st.title("Medical Report Review System")

    # Initialize session state
    initialize_session_state()

    # 获取输入方式
    input_selection = setup_input_selection()
    
    if input_selection:
        if input_selection["type"] == "dataset":
            # 原有的数据集选择逻辑
            selected_report = setup_report_selection()
            if selected_report:
                report_data = st.session_state.reports_json[selected_report]
                entities_data = report_data['entities']
                selections_og = entities2Selection(report_data['text'], entities_data)
        else:
            # 处理用户输入的文本
            user_text = input_selection["text"]
            if user_text:
                from ie import process_sample
                entities_data = process_sample(user_text, num_shots=5)
                report_data = user_text
                selections_og = entities2Selection(user_text, entities_data)
            else:
                st.warning("请输入文本内容")
                return

        # 创建两列布局
        col1, col2 = st.columns([2, 2])  # 调整列宽比例

        # Display report content
        with col1:
            display_report_content(report_data)
            # Display relationship graph
            current_entities = selections_og
            if st.button("Update Graph", key="update_graph"):
                current_entities = selections_og
            display_relationship_graph(current_entities, entities_data)

        # Display entities
        with col2:
            selections = display_entities(
                report_data['text'] if isinstance(report_data, dict) else report_data, 
                selections_og
            )

        # 仅对数据集的报告显示提交按钮
        if input_selection["type"] == "dataset":
            handle_review_submission(selected_report, selections, entities_data)

if __name__ == "__main__":
    main()