import os
import gradio as gr
from langchain_core.prompts import PromptTemplate
from langchain_community.document_loaders import PyPDFLoader
from langchain_google_genai import ChatGoogleGenerativeAI
import google.generativeai as genai
from langchain.chains.question_answering import load_qa_chain
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from PIL import Image
import io
from langchain_groq import ChatGroq
from opencc import OpenCC

# 設置OpenCC轉換器
cc = OpenCC('s2t')  # 從簡體轉換到繁體

# 配置Gemini API
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))

# 配置Groq API
groq_api_key = os.getenv("GROQ_API_KEY")
os.environ["GROQ_API_KEY"] = groq_api_key

# 加載Mistral模型
model_path = "nvidia/Mistral-NeMo-Minitron-8B-Base"
mistral_tokenizer = AutoTokenizer.from_pretrained(model_path)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
dtype = torch.bfloat16
mistral_model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=dtype, device_map=device)

def to_traditional_chinese(text):
    return cc.convert(text)

def process_with_gemini(file, image, question):
    gemini_model = ChatGoogleGenerativeAI(model="gemini-pro", temperature=0.3)
    
    if file:
        pdf_loader = PyPDFLoader(file.name)
        pages = pdf_loader.load_and_split()
        context = "\n".join(str(page.page_content) for page in pages[:30])
        
        prompt_template = """根據提供的上下文盡可能準確地回答問題。如果上下文中沒有答案，請說"上下文中沒有可用的答案" \n\n 上下文: \n {context}?\n 問題: \n {question} \n 回答: """
        prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
        
        chain = load_qa_chain(gemini_model, chain_type="stuff", prompt=prompt)
        result = chain({"input_documents": pages, "question": question, "context": context}, return_only_outputs=True)
        return to_traditional_chinese(result['output_text'])
    elif image:
        vision_model = genai.GenerativeModel('gemini-pro-vision')
        response = vision_model.generate_content([image, question])
        return to_traditional_chinese(response.text)
    else:
        return "請上傳PDF文件或圖片。"

def process_with_groq(file, question):
    if not file:
        return "Groq處理只適用於PDF文件。請上傳PDF文件。"
    
    groq_model = ChatGroq(model_name="mixtral-8x7b-32768", temperature=0.3)
    pdf_loader = PyPDFLoader(file.name)
    pages = pdf_loader.load_and_split()
    context = "\n".join(str(page.page_content) for page in pages[:30])
    
    prompt_template = """根據提供的上下文盡可能準確地回答問題。如果上下文中沒有答案，請說"上下文中沒有可用的答案" \n\n 上下文: \n {context}?\n 問題: \n {question} \n 回答: """
    prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
    
    chain = load_qa_chain(groq_model, chain_type="stuff", prompt=prompt)
    result = chain({"input_documents": pages, "question": question, "context": context}, return_only_outputs=True)
    return to_traditional_chinese(result['output_text'])

def process_with_mistral(question):
    mistral_prompt = f"根據這個問題: {question}\n生成一個回答:"
    mistral_inputs = mistral_tokenizer.encode(mistral_prompt, return_tensors='pt').to(device)
    with torch.no_grad():
        mistral_outputs = mistral_model.generate(mistral_inputs, max_length=100)
    mistral_output = mistral_tokenizer.decode(mistral_outputs[0], skip_special_tokens=True)
    return to_traditional_chinese(mistral_output)

# 定義Gradio界面
with gr.Blocks() as demo:
    gr.Markdown("# 多模態RAG知識檢索系統（使用Gemini、Groq和Mistral）")
    
    with gr.Row():
        input_file = gr.File(label="上傳PDF文件")
        input_image = gr.Image(type="pil", label="上傳圖片")
    
    input_question = gr.Textbox(label="詢問文檔或圖片相關問題")
    
    with gr.Row():
        gemini_button = gr.Button("使用Gemini處理")
        groq_button = gr.Button("使用Groq處理")
        mistral_button = gr.Button("使用Mistral處理")
    
    with gr.Row():
        gemini_output = gr.Textbox(label="Gemini輸出")
        groq_output = gr.Textbox(label="Groq輸出")
        mistral_output = gr.Textbox(label="Mistral輸出")
    
    gemini_button.click(
        fn=process_with_gemini,
        inputs=[input_file, input_image, input_question],
        outputs=gemini_output
    )
    
    groq_button.click(
        fn=process_with_groq,
        inputs=[input_file, input_question],
        outputs=groq_output
    )
    
    mistral_button.click(
        fn=process_with_mistral,
        inputs=[input_question],
        outputs=mistral_output
    )

demo.launch()