from transformers import Qwen2VLForConditionalGeneration, AutoProcessor

model = Qwen2VLForConditionalGeneration.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")

import streamlit as st
import os
from byaldi import RAGMultiModalModel
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
from pdf2image import convert_from_path
from streamlit_option_menu import option_menu
from datetime import datetime

# 設置頁面配置
st.set_page_config(page_title="多模態RAG系統", layout="wide")

# 創建必要的目錄
UPLOAD_DIR = "uploaded_pdfs"
if not os.path.exists(UPLOAD_DIR):
    os.makedirs(UPLOAD_DIR)

@st.cache_resource
def load_models():
    with st.spinner("正在加載模型... 這可能需要幾分鐘。"):
        rag_engine = RAGMultiModalModel.from_pretrained("vidore/colpali")
        model = Qwen2VLForConditionalGeneration.from_pretrained("Qwen/Qwen2-VL-7B-Instruct", torch_dtype=torch.float16, device_map="cuda")

        processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct", trust_remote_code=True)
    return rag_engine, model, processor

def save_uploaded_file(uploaded_file):
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"{timestamp}_{uploaded_file.name}"
    file_path = os.path.join(UPLOAD_DIR, filename)
    with open(file_path, "wb") as f:
        f.write(uploaded_file.getvalue())
    return file_path

def main():
    if 'indexed_files' not in st.session_state:
        st.session_state.indexed_files = set()

    selected = option_menu(menu_title=None, optinotallow=["上傳PDF", "查詢文檔"], icnotallow=["cloud-upload", "search"], default_index=0)

    rag_engine, model, processor = load_models()

    if selected == "上傳PDF":
        st.title("PDF文檔上傳")
        uploaded_files = st.file_uploader("上傳你的PDF文檔", type=['pdf'], accept_multiple_files=True)

        if uploaded_files:
            for uploaded_file in uploaded_files:
                if uploaded_file.name not in [os.path.basename(f) for f in st.session_state.indexed_files]:
                    with st.spinner(f"正在處理{uploaded_file.name}..."):
                        file_path = save_uploaded_file(uploaded_file)
                        try:
                            rag_engine.index(input_path=file_path, index_name=os.path.basename(file_path), store_collection_with_index=True, overwrite=True)
                            st.session_state.indexed_files.add(file_path)
                            st.success(f"成功處理{uploaded_file.name}")
                        except Exception as e:
                            st.error(f"處理{uploaded_file.name}時出錯：{str(e)}")

elif selected == "查詢文檔":
    st.title("查詢文檔")

    if not st.session_state.indexed_files:
        st.warning("請先上傳并索引一些文檔！")
        return

    query = st.text_input("輸入你的查詢：", placeholder="你想知道什么？")

    if query:
        with st.spinner("正在處理查詢..."):
            all_results = []
            for file_path in st.session_state.indexed_files:
                results = rag_engine.search(query, k=3, index_name=os.path.basename(file_path))
                all_results.extend([(file_path, r) for r in results])

            all_results.sort(key=lambda x: x[1].get('score', 0), reverse=True)

            if all_results:
                top_file, top_result = all_results[0]
                images = convert_from_path(top_file)
                image_index = top_result["page_num"] - 1

                # 在標簽頁中顯示結果
                tab1, tab2 = st.tabs(["結果", "上下文"])

                with tab1:
                    col1, col2 = st.columns([1, 1])
                    with col1:
                        st.image(images[image_index], captinotallow=f"來自{os.path.basename(top_file)}的第{image_index + 1}頁", use_column_width=True)
                    with col2:
                        messages = [{"role": "user", "content": [{"type": "image", "image": images[image_index]}, {"type": "text", "text": query}]}]
                        text = processor.apply_chat_template(messages)
                        inputs = processor(text=[text], images=[images[image_index]], padding=True, return_tensors="pt").to("cuda")

                        generated_ids = model.generate(**inputs, max_new_tokens=50)
                        output_text = processor.batch_decode(generated_ids[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)

                        st.markdown("### 模型響應")
                        st.write(output_text[0])

                with tab2:
                    for file_path, result in all_results[:5]:
                        with st.expander(f"來自：{os.path.basename(file_path)} - 第{result['page_num']}頁"):
                            st.write(result["content"])
                            st.caption(f"相關性得分：{result.get('score', 0):.2f}")