import gradio as gr
import pandas as pd
from datasets import load_dataset
from googletrans import Translator
from huggingface_hub import login
import os

# --- 定数と初期設定 ---
# Hugging Faceトークンを環境変数から取得
HF_TOKEN = os.environ.get("HF_TOKEN")

# --- グローバル変数 ---
# データは一度だけロードする
dataset = None
df = None
category_counts = None

# --- データ処理関数 ---
def load_and_process_data():
    """データセットをロードし、グローバル変数を初期化する"""
    global dataset, df, category_counts
    if dataset is not None:
        return

    try:
        # トークンがあればログイン
        if HF_TOKEN:
            login(token=HF_TOKEN)
        dataset = load_dataset("cais/hle", split="test")
        # 画像データを扱うため、Pandasに変換するのはメタデータのみ
        df = dataset.remove_columns(['image_preview', 'rationale_image']).to_pandas()
        category_counts = df['category'].value_counts()
        print("データセットのロードと前処理が完了しました。")
    except Exception as e:
        print(f"データセットのロードエラー: {e}")
        # エラーが発生した場合、アプリがクラッシュしないように空のデータフレームを設定
        df = pd.DataFrame(columns=['id', 'question', 'category'])
        category_counts = pd.Series()

# --- 翻訳関数 ---
translator = Translator()
def translate_text(text, dest_lang='ja'):
    if not text or not isinstance(text, str):
        return ""
    try:
        # Translatorインスタンスを再利用
        return translator.translate(text, dest=dest_lang).text
    except Exception as e:
        return f"翻訳エラー: {e}"

# --- Gradioイベントハンドラ ---
def on_category_change(selected_category):
    """カテゴリが変更されたときに、問題のドロップダウンを更新する"""
    if selected_category == "全カテゴリ":
        filtered_indices = df.index
    else:
        filtered_indices = df[df['category'] == selected_category].index

    # 選択肢を (表示名, 値) のタプル形式で作成
    question_choices = [
        (f"{df.loc[idx, 'question'][:80]}...", idx) for idx in filtered_indices
    ]

    if not question_choices:
        # 選択肢がない場合は、表示をクリアし、選択不可にする
        return gr.Dropdown(choices=[], label="問題 (該当なし)", interactive=False, value=None)
    else:
        return gr.Dropdown(choices=question_choices, label="問題を選択", interactive=True, value=question_choices[0][1])

def on_question_change(selected_index):
    """問題が選択されたときに、すべての詳細表示を更新する"""
    if selected_index is None or pd.isna(selected_index):
        # 空の出力をまとめて返す
        empty_outputs = [gr.Markdown(visible=False)] * 6 + [gr.Image(visible=False)] * 2
        return tuple(empty_outputs)

    # 元のHugging Face Datasetから完全なエントリを取得
    entry = dataset[int(selected_index)]

    # 各要素の翻訳
    q_trans = translate_text(entry['question'])
    a_trans = translate_text(entry['answer'])
    r_trans = translate_text(entry.get('rationale', ''))

    # 出力コンポーネントの値を生成
    outputs = {
        "question_md": gr.Markdown(f"### 質問\n---\n**原文:**\n{entry['question']}\n\n**日本語訳:**\n{q_trans}", visible=True),
        "question_img": gr.Image(entry.get('image_preview'), label="質問画像", visible=bool(entry.get('image_preview'))),
        "answer_md": gr.Markdown(f"### 回答\n---\n**原文:**\n{entry['answer']}\n\n**日本語訳:**\n{a_trans}", visible=True),
        "rationale_md": gr.Markdown(f"### 解説\n---\n**原文:**\n{entry.get('rationale', 'N/A')}\n\n**日本語訳:**\n{r_trans}", visible=bool(entry.get('rationale'))),
        "rationale_img": gr.Image(entry.get('rationale_image'), label="解説画像", visible=bool(entry.get('rationale_image'))),
        "metadata_md": gr.Markdown(f"**ID:** `{entry['id']}`<br>**分野:** `{entry['raw_subject']}`<br>**回答タイプ:** `{entry['answer_type']}`", visible=True),
        "json_output": gr.JSON({k: str(v) for k, v in entry.items()}, label="元のデータ", visible=True)
    }

    # 定義された順序で値を返す
    return (
        outputs["question_md"],
        outputs["question_img"],
        outputs["answer_md"],
        outputs["rationale_md"],
        outputs["rationale_img"],
        outputs["metadata_md"],
        outputs["json_output"]
    )

# --- Gradio UI構築 ---
def create_demo():
    # アプリケーション起動時に一度だけデータをロード
    load_and_process_data()

    with gr.Blocks(theme=gr.themes.Soft(), title="HLE Dataset Viewer") as demo:
        gr.Markdown("# Humanity's Last Exam (HLE) Dataset Viewer")
        gr.Markdown("Hugging Face `cais/hle`データセットを探索し、日本語訳を確認できます。")

        with gr.Row():
            with gr.Column(scale=1, min_width=350):
                gr.Markdown("## 操作パネル")
                category_dd = gr.Dropdown(
                    choices=["全カテゴリ"] + sorted(category_counts.index.tolist()),
                    value="全カテゴリ",
                    label="1. カテゴリを選択"
                )
                question_dd = gr.Dropdown(label="2. 問題を選択", interactive=False)

                gr.Markdown("### カテゴリ別問題数")
                gr.Dataframe(value=pd.DataFrame(category_counts).reset_index(), headers=['カテゴリ', '問題数'], interactive=False)

            with gr.Column(scale=3):
                # 出力エリアのプレースホルダー
                metadata_md = gr.Markdown(visible=False)
                question_md = gr.Markdown(visible=False)
                question_img = gr.Image(label="質問画像", visible=False)
                answer_md = gr.Markdown(visible=False)
                rationale_md = gr.Markdown(visible=False)
                rationale_img = gr.Image(label="解説画像", visible=False)
                json_output = gr.JSON(label="元のデータ", visible=False)

        # イベントリスナーを設定
        category_dd.change(fn=on_category_change, inputs=category_dd, outputs=question_dd)
        question_dd.change(fn=on_question_change, inputs=question_dd, outputs=[
            question_md, question_img, answer_md, rationale_md, rationale_img, metadata_md, json_output
        ])

        # 初期表示のために最初のカテゴリ変更イベントをトリガー
        demo.load(fn=on_category_change, inputs=category_dd, outputs=question_dd)

    return demo

# --- アプリケーション起動 ---
if __name__ == "__main__":
    app = create_demo()
    app.launch()