import gradio as gr
import os
import time
import pdfplumber
from dotenv import load_dotenv
import torch
from transformers import (
    BertJapaneseTokenizer,
    BertModel,
    AutoTokenizer,
    AutoModelForCausalLM,
    pipeline,
    BitsAndBytesConfig
)
from langchain_community.vectorstores import FAISS  # 修正
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
from langchain_community.llms import HuggingFacePipeline  # 修正
from langchain_community.embeddings import HuggingFaceEmbeddings  # 修正
from langchain_huggingface import HuggingFaceEndpoint

# Pydanticの警告を無視
import warnings
warnings.filterwarnings(
    "ignore",
    message=r"Field \"model_name\" in HuggingFaceInferenceAPIEmbeddings has conflict with protected namespace"
)

load_dotenv()

list_llm = [
    "meta-llama/Meta-Llama-3-8B-Instruct",
    "rinna/llama-3-youko-8b",
]
list_llm_simple = [os.path.basename(llm) for llm in list_llm]

# 日本語PDFのテキスト抽出
def extract_text_from_pdf(file_path):
    with pdfplumber.open(file_path) as pdf:
        pages = [page.extract_text() for page in pdf.pages]
    return " ".join(pages)

# モデルとトークナイザの初期化
tokenizer_bert = BertJapaneseTokenizer.from_pretrained(
    'cl-tohoku/bert-base-japanese',
    clean_up_tokenization_spaces=True
)
model_bert = BertModel.from_pretrained('cl-tohoku/bert-base-japanese')

def split_text_simple(text, chunk_size=1024):
    return [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]

def create_db(splits):
    embeddings = HuggingFaceEmbeddings(
        model_name='sonoisa/sentence-bert-base-ja-mean-tokens'
    )
    vectordb = FAISS.from_texts(splits, embeddings)
    return vectordb

def initialize_llmchain(
    llm_model,
    temperature,
    max_tokens,
    top_k,
    vector_db,
    retries=5,
    delay=5
):
    attempt = 0
    while attempt < retries:
        try:
            # ローカルモデルの場合
            if "rinna" in llm_model.lower():
                # デバイスの自動検出
                if torch.cuda.is_available():
                    device_map = "auto"
                    torch_dtype = torch.float16
                    # GPUがある場合は量子化を使用
                    quantization_config = BitsAndBytesConfig(
                        load_in_4bit=True,
                        bnb_4bit_compute_dtype=torch.float16,
                        bnb_4bit_use_double_quant=True,
                        bnb_4bit_quant_type="nf4"
                    )
                    model = AutoModelForCausalLM.from_pretrained(
                        llm_model,
                        device_map=device_map,
                        quantization_config=quantization_config
                    )
                else:
                    device_map = {"": "cpu"}
                    torch_dtype = torch.float32
                    # CPUの場合は量子化を使用せずにモデルをロード
                    model = AutoModelForCausalLM.from_pretrained(
                        llm_model,
                        device_map=device_map,
                        torch_dtype=torch_dtype
                    )
                tokenizer = AutoTokenizer.from_pretrained(llm_model, use_fast=False)
                pipe = pipeline(
                    "text-generation",
                    model=model,
                    tokenizer=tokenizer,
                    max_new_tokens=max_tokens,
                    temperature=temperature
                )
                llm = HuggingFacePipeline(pipeline=pipe)
            # エンドポイントモデルの場合
            elif "meta-llama" in llm_model.lower() or "mistralai" in llm_model.lower():
                # パラメータを直接指定
                llm = HuggingFaceEndpoint(
                    endpoint_url=f"https://api-inference.huggingface.co/models/{llm_model}",
                    huggingfacehub_api_token=os.getenv("HF_TOKEN"),
                    temperature=temperature,
                    max_new_tokens=max_tokens,
                    top_k=top_k
                )
            else:
                # その他のモデルの場合（必要に応じて追加）
                raise Exception(f"Unsupported model: {llm_model}")

            # 共通の処理
            memory = ConversationBufferMemory(
                memory_key="chat_history",
                output_key='answer',
                return_messages=True
            )
            retriever = vector_db.as_retriever()
            qa_chain = ConversationalRetrievalChain.from_llm(
                llm,
                retriever=retriever,
                memory=memory,
                return_source_documents=True,
                verbose=False
            )
            return qa_chain
        except Exception as e:
            if "Could not authenticate with huggingface_hub" in str(e):
                time.sleep(delay)
                attempt += 1
            else:
                raise Exception(f"Error initializing QA chain: {str(e)}")
    raise Exception(f"Failed to initialize after {retries} attempts")

def process_pdf(file):
    try:
        if file is None:
            return None, "Please upload a PDF file."
        text = extract_text_from_pdf(file.name)
        splits = split_text_simple(text)
        vdb = create_db(splits)
        return vdb, "PDF processed and vector database created."
    except Exception as e:
        return None, f"Error processing PDF: {str(e)}"

def initialize_qa_chain(
    llm_index,
    temperature,
    max_tokens,
    top_k,
    vector_db
):
    try:
        if vector_db is None:
            return None, "Please process a PDF first."
        llm_name = list_llm[llm_index]
        chain = initialize_llmchain(
            llm_name,
            temperature,
            max_tokens,
            top_k,
            vector_db
        )
        return chain, "QA Chatbot initialized with selected LLM."
    except Exception as e:
        return None, f"Error initializing QA chain: {str(e)}"

def update_chat(msg, history, chain):
    try:
        if chain is None:
            return history + [("User", msg), ("Assistant", "Please initialize the QA Chatbot first.")]
        response = chain({"question": msg, "chat_history": history})
        return history + [("User", msg), ("Assistant", response['answer'])]
    except Exception as e:
        return history + [("User", msg), ("Assistant", f"Error: {str(e)}")]

def demo():
    with gr.Blocks() as demo:
        vector_db = gr.State(value=None)
        qa_chain = gr.State(value=None)

        with gr.Tab("Step 1 - Upload and Process"):
            with gr.Row():
                document = gr.File(label="Upload your Japanese PDF document", file_types=["pdf"])
            with gr.Row():
                process_btn = gr.Button("Process PDF")
                process_output = gr.Textbox(label="Processing Output")

        with gr.Tab("Step 2 - Initialize QA Chatbot"):
            with gr.Row():
                llm_btn = gr.Radio(list_llm_simple, label="Select LLM Model", type="index")
                llm_temperature = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, label="Temperature", value=0.7)
                max_tokens = gr.Slider(minimum=128, maximum=2048, step=128, label="Max Tokens", value=1024)
                top_k = gr.Slider(minimum=1, maximum=10, step=1, label="Top K", value=3)
            with gr.Row():
                init_qa_btn = gr.Button("Initialize QA Chatbot")
                init_output = gr.Textbox(label="Initialization Output")

        with gr.Tab("Step 3 - Chat with your Document"):
            chatbot = gr.Chatbot()
            message = gr.Textbox(label="Ask a question")
            with gr.Row():
                send_btn = gr.Button("Send")
                clear_chat_btn = gr.Button("Clear Chat")
                reset_all_btn = gr.Button("Reset All")

        process_btn.click(
            process_pdf,
            inputs=[document],
            outputs=[vector_db, process_output]
        )

        init_qa_btn.click(
            initialize_qa_chain,
            inputs=[llm_btn, llm_temperature, max_tokens, top_k, vector_db],
            outputs=[qa_chain, init_output]
        )

        send_btn.click(
            update_chat,
            inputs=[message, chatbot, qa_chain],
            outputs=[chatbot]
        )

        # Clear Chatボタン：チャット履歴のみをクリア
        clear_chat_btn.click(
            lambda: None,
            outputs=[chatbot]
        )

        # Reset Allボタン：チャット履歴、PDFデータ、チャットボットの状態をすべてクリア
        reset_all_btn.click(
            lambda: (None, None, None),
            outputs=[chatbot, vector_db, qa_chain]
        )
    return demo

if __name__ == "__main__":
    demo().launch()