import os

import gradio as gr
from transformers.utils import logging

from langchain_ollama import OllamaEmbeddings
from langchain_community.vectorstores import Neo4jVector

logging.set_verbosity_info()
logger = logging.get_logger("transformers")

# Neo4jへの接続情報
NEO4J_URL = os.environ['NEO4J_URL']
NEO4J_USERNAME = os.environ['NEO4J_USERNAME']
NEO4J_PASSWORD = os.environ['NEO4J_PASSWORD']
NEO4J_DATABASE = os.environ['NEO4J_DATABASE']

EMBEDDINGS = OllamaEmbeddings(
    model="mxbai-embed-large",
)

def hybrid_search(input_text, top_k):

    # グラフからノード検索用インデックスを取得
    index = Neo4jVector.from_existing_graph(
        embedding=EMBEDDINGS,
        url=NEO4J_URL,
        username=NEO4J_USERNAME,
        password=NEO4J_PASSWORD,
        database=NEO4J_DATABASE,
        node_label="Document", # 検索対象ノード
        text_node_properties=["id", "text"], # 検索対象プロパティ
        embedding_node_property="embedding", # ベクトルデータの保存先プロパティ
        index_name="vector_index", # ベクトル検索用のインデックス名
        keyword_index_name="fulltext_index", # 全文検索用のインデックス名
        search_type="hybrid" # 検索タイプに「ハイブリッド」を設定（デフォルトは「ベクター」）
    )

    all_answers = []
    # クエリを設定して検索を実行
    query = input_text
    docs_with_score = index.similarity_search_with_score(query, k=top_k)
    for i in docs_with_score:
        doc, score = i
        all_answers.append(doc.metadata["source"])

    return "\n***\n".join(all_answers)


CSS ="""
.contain { display: flex; flex-direction: column; }
.gradio-container { height: 100vh !important; }
#component-0 { height: 100%; }
#textbox { flex-grow: 1; overflow: auto; resize: vertical; }
.secondary {background-color: #6366f1; }
#full-width-button { width: 100%; }
"""
#with gr.Blocks() as demo:
with gr.Blocks(theme=gr.themes.Monochrome(radius_size=gr.themes.sizes.radius_sm)) as demo:
    with gr.Row():
        gr.Markdown("# 裁定検索")
    with gr.Row():
        output = gr.TextArea(
            elem_id="検索結果",
            label="検索結果",
        )
    with gr.Row():
        input = gr.Textbox(
            label="質問",
            placeholder="芸魔龍王アメイジンの出た時の効果は、後から出たクリーチャーも影響しますか",
            lines=3,
        )
    with gr.Row():
        submit = gr.Button(value="検索", variant="secondary", elem_id="full-width-button")
        top_k = gr.Slider(1, 10, label="表示数", step=1, value=5, interactive=True)

    submit_click_event = submit.click(fn=hybrid_search, inputs=[input, top_k], outputs=output)

demo.launch()