import os

import gradio as gr
from gradio.themes.base import Base
from gradio.themes.utils import colors, fonts, sizes
from transformers.utils import logging

from langchain_ollama import OllamaEmbeddings
from langchain_huggingface import HuggingFaceEmbeddings
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",
# )

EMBEDDINGS = HuggingFaceEmbeddings(
    model_name="mixedbread-ai/mxbai-embed-large-v1"
)


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************************************************************\n\n".join(all_answers)


# カスタムテーマの作成
class CustomTheme(Base):
    def __init__(
        self,
        *,
        primary_hue: colors.Color | str = colors.blue,
        secondary_hue: colors.Color | str = colors.gray,
        neutral_hue: colors.Color | str = colors.gray,
        spacing_size: sizes.Size | str = sizes.spacing_md,
        radius_size: sizes.Size | str = sizes.radius_md,
        text_size: sizes.Size | str = sizes.text_lg,  # フォントサイズを大きく設定
        font: fonts.Font
        | str
        | list[fonts.Font | str] = (
            fonts.GoogleFont("IBM Plex Sans"),
            "Arial",
            "sans-serif",
        ),
        font_mono: fonts.Font
        | str
        | list[fonts.Font | str] = (
            fonts.GoogleFont("IBM Plex Mono"),
            "Courier",
            "monospace",
        ),
    ):
        super().__init__(
            primary_hue=primary_hue,
            secondary_hue=secondary_hue,
            neutral_hue=neutral_hue,
            spacing_size=spacing_size,
            radius_size=radius_size,
            text_size=text_size,
            font=font,
            font_mono=font_mono,
        )

# カスタムテーマの適用
theme = CustomTheme()

CSS ="""
.contain { display: flex; flex-direction: column; }
.gradio-container { min-height: 100vh !important; }
#component-0 { height: 100%; }
#textbox { flex-grow: 1; overflow: auto; resize: vertical;}
.secondary {background-color: #6366f1; }
#full-width-button { width: 100%; }
#search-result { overflow-y: scroll !important; font-size:18px !important; font-weight:500 !important;}
#question-box { font-size:18px !important; font-weight:500 !important; }
"""
#with gr.Blocks() as demo:
with gr.Blocks(theme=theme, css=CSS) as demo:
    with gr.Row():
        with gr.Column():
            gr.Markdown(f"""
                    ### ・非公式サイトです  
                    ### ・デモでしかないので速度・精度・動作は保証しないし新しい裁定にも対応しません。突然消す可能性もあり
                    ### ・ですます調で質問をすると精度が上がるかも\n\n
                    """)
            
    # スペースを追加
    with gr.Row():
        gr.Markdown("<br><br>", elem_id="spacer", visible=True)  # 改行タグを挿入してスペースを作成
    
    with gr.Row():
        gr.Markdown("# 裁定検索")
    with gr.Row():
        output = gr.TextArea(
            elem_id="search-result",
            label="検索結果",
        )
    with gr.Row():
        input = gr.Textbox(
            elem_id="question-box",
            label="質問",
            placeholder="カードを指定して破壊する能力でクリーチャーの下にあるカードを指定できますか",
            lines=3,
        )
    with gr.Row():
        submit = gr.Button(value="検索", variant="huggingface", 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()